From 4157a20cec17069de2da1a46d1ca6f25ef890a70 Mon Sep 17 00:00:00 2001 From: thomasloux Date: Thu, 18 Sep 2025 13:57:07 +0000 Subject: [PATCH 1/2] fix:orb squeeze incorrect energy shape --- torch_sim/models/orb.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch_sim/models/orb.py b/torch_sim/models/orb.py index fd65b23f..132f6d5c 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] From 780fc23bfa05376e05a2571e80a4b040a277db40 Mon Sep 17 00:00:00 2001 From: thomasloux Date: Thu, 13 Nov 2025 18:50:49 +0100 Subject: [PATCH 2/2] fix cell_to_cellpar --- torch_sim/models/orb.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch_sim/models/orb.py b/torch_sim/models/orb.py index 7fda6c1f..b998c7a5 100644 --- a/torch_sim/models/orb.py +++ b/torch_sim/models/orb.py @@ -82,7 +82,7 @@ def cell_to_cellpar( x = torch.dot(cell[j], cell[k]) / ll angle = 180.0 / torch.pi * torch.arccos(x) else: - angle = 90.0 + angle = torch.tensor(90.0, device=cell.device, dtype=cell.dtype) angles.append(angle) if radians: angles = [angle * torch.pi / 180 for angle in angles]