From ab32e94a030a48469a21a90de7726a79d7ddd924 Mon Sep 17 00:00:00 2001
From: Ilyes Batatia <48651863+ilyes319@users.noreply.github.com>
Date: Wed, 10 Jul 2024 21:18:54 +0100
Subject: [PATCH] change dynamo to optional

---
 mace/tools/compile.py | 5 ++++-
 tests/test_hessian.py | 3 ++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/mace/tools/compile.py b/mace/tools/compile.py
index 1832e312..425e4c02 100644
--- 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
index 0b1c5d86..53457335 100644
--- 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