optional import from experimental

FAIR-Chem · Aug 20, 2024 · f3e1c38 · f3e1c38
1 parent 21eecd4
commit f3e1c38
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/src/fairchem/core/datasets/ase_datasets.py b/src/fairchem/core/datasets/ase_datasets.py
@@ -29,11 +29,9 @@
 from fairchem.core.datasets.lmdb_database import LMDBDatabase
 from fairchem.core.datasets.target_metadata_guesser import guess_property_metadata
 
-# from fairchem.core.modules.transforms import DataTransforms
 from fairchem.core.preprocessing import AtomsToGraphs
-from fairchem.experimental.foundation_models.multi_task_dataloader.transforms.data_object import (
-    DataTransforms,
-)
+
+from fairchem.core.modules.transforms import DataTransforms
 
 
 def apply_one_tags(

diff --git a/src/fairchem/core/modules/transforms.py b/src/fairchem/core/modules/transforms.py
@@ -9,6 +9,11 @@
 if TYPE_CHECKING:
     from torch_geometric.data import Data
 
+try:
+    from fairchem.experimental.foundation_models.multi_task_dataloader.transforms.data_object import *
+except ImportError:
+    pass
+
 
 class DataTransforms:
     def __init__(self, config) -> None: