avoid hf loading tasks

atomind-ai · Jan 13, 2025 · a909029 · a909029
1 parent 7ab307c
commit a909029
Showing 1 changed file with 10 additions and 8 deletions.
diff --git a/mlip_arena/tasks/__init__.py b/mlip_arena/tasks/__init__.py
@@ -6,15 +6,17 @@
 from mlip_arena.models import MLIP
 from mlip_arena.models import REGISTRY as MODEL_REGISTRY
 
+try:
+    from .elasticity import run as ELASTICITY
+    from .eos import run as EOS
+    from .md import run as MD
+    from .neb import run as NEB
+    from .neb import run_from_endpoints as NEB_FROM_ENDPOINTS
+    from .optimize import run as OPT
 
-from .elasticity import run as ELASTICITY
-from .eos import run as EOS
-from .md import run as MD
-from .neb import run as NEB
-from .neb import run_from_endpoints as NEB_FROM_ENDPOINTS
-from .optimize import run as OPT
-
-__all__ = ["OPT", "EOS", "MD", "NEB", "NEB_FROM_ENDPOINTS", "ELASTICITY"]
+    __all__ = ["OPT", "EOS", "MD", "NEB", "NEB_FROM_ENDPOINTS", "ELASTICITY"]
+except ImportError:
+    pass
 
 with open(Path(__file__).parent / "registry.yaml", encoding="utf-8") as f:
     REGISTRY = yaml.safe_load(f)