diff --git a/hippynn/molecular_dynamics/md.py b/hippynn/molecular_dynamics/md.py
index dab80a53..edb76b8c 100644
--- a/hippynn/molecular_dynamics/md.py
+++ b/hippynn/molecular_dynamics/md.py
@@ -458,14 +458,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.