fix bug with batched MD saving

shinkle-lanl · Oct 17, 2024 · ebc2221 · ebc2221
1 parent 144c160
commit ebc2221
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/hippynn/molecular_dynamics/__init__.py b/hippynn/molecular_dynamics/__init__.py
@@ -5,7 +5,7 @@
 
 """
 
-from .md import MolecularDynamics, Variable, NullUpdater, VelocityVerlet, LangevinDynamics
+from .md import MolecularDynamics, Variable, NullUpdater, VelocityVerlet, LangevinDynamics, VariableUpdater
 
 
-__all__ = ["MolecularDynamics", "Variable", "NullUpdater", "VelocityVerlet", "LangevinDynamics"]
+__all__ = ["MolecularDynamics", "Variable", "VariableUpdater", "NullUpdater", "VelocityVerlet", "LangevinDynamics"]
diff --git a/hippynn/molecular_dynamics/md.py b/hippynn/molecular_dynamics/md.py
@@ -450,14 +450,14 @@ def _update_data(self, model_outputs: dict):
         for variable in self.variables:
             for key, value in variable.data.items():
                 try:
-                    self._data[f"{variable.name}_{key}"].append(value.cpu().detach()[0])
+                    self._data[f"{variable.name}_{key}"].append(value.cpu().detach())
                 except KeyError:
-                    self._data[f"{variable.name}_{key}"] = [value.cpu().detach()[0]]
+                    self._data[f"{variable.name}_{key}"] = [value.cpu().detach()]
         for key, value in model_outputs.items():
             try:
-                self._data[f"output_{key}"].append(value.cpu().detach()[0])
+                self._data[f"output_{key}"].append(value.cpu().detach())
             except KeyError:
-                self._data[f"output_{key}"] = [value.cpu().detach()[0]]
+                self._data[f"output_{key}"] = [value.cpu().detach()]
 
     def run(self, dt: float, n_steps: int, record_every: Optional[int] = None):
         """Run `n_steps` of MD algorithm.