fixed missing perturbation arg for ensemble in jaxmd

apax-hub · Jan 20, 2024 · 41884c3 · 41884c3
1 parent 4159f6c
commit 41884c3
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/apax/md/nvt.py b/apax/md/nvt.py
@@ -30,9 +30,9 @@
 
 
 def create_energy_fn(model, params, numbers, n_models):
-    def ensemble(params, R, Z, neighbor, box, offsets):
-        vmodel = jax.vmap(model, (0, None, None, None, None, None), 0)
-        energies = vmodel(params, R, Z, neighbor, box, offsets)
+    def ensemble(params, R, Z, neighbor, box, offsets, perturbation=None):
+        vmodel = jax.vmap(model, (0, None, None, None, None, None, None), 0)
+        energies = vmodel(params, R, Z, neighbor, box, offsets, perturbation)
         energy = jnp.mean(energies)
 
         return energy