necessary when the graph is empty.

ACEsuit · Nov 28, 2022 · 5a69c24 · 5a69c24
1 parent f49b10a
commit 5a69c24
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/mace/calculators/lammps_mace.py b/mace/calculators/lammps_mace.py
@@ -70,6 +70,8 @@ def forward(
                     torch.cross(cell[:, 1, :], cell[:, 2, :], dim=1),
                 ).unsqueeze(-1)
                 stress = virials / volume.view(-1, 1, 1)
+            else:
+                virials = torch.zeros_like(displacement)
 
         total_energy = scatter_sum(
             src=node_energy, index=data["batch"], dim=-1, dim_size=num_graphs