remove id_tab

thorben-frank · Aug 23, 2024 · 02107c2 · 02107c2
1 parent e8d59ae
commit 02107c2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mlff/mdx/simulate.py b/mlff/mdx/simulate.py
@@ -6,7 +6,7 @@
 import logging
 
 from jax import lax
-from jax.experimental.host_callback import id_tap
+# from jax.experimental.host_callback import id_tap
 
 from typing import Dict, Tuple
 from collections import namedtuple
@@ -116,7 +116,7 @@ def run(self, integrator, atoms, steps):
             integrator = final_integrator
 
             if step % self.save_frequency == 0:
-                id_tap(self.save_trajectory, step_data)
+                self.save_trajectory(step_data, None)
 
             t = delta_t * (step + 1 - overflow_count)
             temp = step_data.observables['temperature'][-1].item()