Make available imports in graphnet.data.{sqlite,parquet} conditional …

…on whether 'torch' is installed
graphnet-team · Aug 30, 2022 · 3ddce57 · 3ddce57
1 parent 37a5a38
commit 3ddce57
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/src/graphnet/data/parquet/__init__.py b/src/graphnet/data/parquet/__init__.py
@@ -1,2 +1,8 @@
+from graphnet.utilities.imports import has_torch_package
+
 from .parquet_dataconverter import ParquetDataConverter
-from .parquet_dataset import ParquetDataset
+
+if has_torch_package():
+    from .parquet_dataset import ParquetDataset
+
+del has_torch_package
diff --git a/src/graphnet/data/sqlite/__init__.py b/src/graphnet/data/sqlite/__init__.py
@@ -1,4 +1,10 @@
+from graphnet.utilities.imports import has_torch_package
+
 from .sqlite_dataconverter import SQLiteDataConverter
-from .sqlite_dataset import SQLiteDataset
-from .sqlite_dataset_perturbed import SQLiteDatasetPerturbed
 from .sqlite_utilities import run_sql_code, save_to_sql
+
+if has_torch_package():
+    from .sqlite_dataset import SQLiteDataset
+    from .sqlite_dataset_perturbed import SQLiteDatasetPerturbed
+
+del has_torch_package