Reenable forces for jax. Fixed little typo.

aimat-lab · Feb 24, 2024 · 2833835 · 2833835
1 parent c54858f
commit 2833835
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/kgcnn/models/force.py b/kgcnn/models/force.py
@@ -117,8 +117,8 @@ def __init__(self,
             self._call_grad_backend = self._call_grad_tf
         elif backend() == "torch":
             self._call_grad_backend = self._call_grad_torch
-        # elif backend() == "jax":
-        #     self._call_grad_backend = self._call_grad_jax
+        elif backend() == "jax":
+            self._call_grad_backend = self._call_grad_jax
         else:
             raise NotImplementedError("Backend '%s' not supported for force model." % backend())