change dynamo to optional

ACEsuit · Jul 10, 2024 · ab32e94 · ab32e94
1 parent 25e3596
commit ab32e94
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/mace/tools/compile.py b/mace/tools/compile.py
@@ -2,7 +2,10 @@
 from functools import wraps
 from typing import Callable, Tuple
 
-import torch._dynamo as dynamo
+try:
+    import torch._dynamo as dynamo
+except ImportError:
+    dynamo = None
 from e3nn import get_optimization_defaults, set_optimization_defaults
 from torch import autograd, nn
 from torch.fx import symbolic_trace

diff --git a/tests/test_hessian.py b/tests/test_hessian.py
@@ -1,6 +1,7 @@
-import pytest
 import numpy as np
+import pytest
 from ase.build import fcc111
+
 from mace.calculators import mace_mp