From 3ddce57c6b86e4f2c664d2775700eeb078a168fc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andreas=20S=C3=B8gaard?= Date: Tue, 30 Aug 2022 21:55:23 +0200 Subject: [PATCH] Make available imports in graphnet.data.{sqlite,parquet} conditional on whether 'torch' is installed --- src/graphnet/data/parquet/__init__.py | 8 +++++++- src/graphnet/data/sqlite/__init__.py | 10 ++++++++-- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/src/graphnet/data/parquet/__init__.py b/src/graphnet/data/parquet/__init__.py index 94f52fc52..05dfbe0c4 100644 --- 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 index ac461ff65..9f1d84d23 100644 --- 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