diff --git a/torch_sim/models/orb.py b/torch_sim/models/orb.py
index acc6f4a1..74364386 100644
--- a/torch_sim/models/orb.py
+++ b/torch_sim/models/orb.py
@@ -416,7 +416,7 @@ def forward(self, state: ts.SimState | StateDict) -> dict[str, torch.Tensor]:
             if not model_has_direct_heads and prop == "stress":
                 continue
             _property = "energy" if prop == "free_energy" else prop
-            results[prop] = predictions[_property].squeeze()
+            results[prop] = predictions[_property]
 
         if self.conservative:
             results["forces"] = results[self.model.grad_forces_name]