remove n_atoms factor

ACEsuit · Jun 5, 2024 · a7b32da · a7b32da
1 parent 0842e7c
commit a7b32da
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/mace/modules/loss.py b/mace/modules/loss.py
@@ -31,11 +31,10 @@ def weighted_mean_squared_stress(ref: Batch, pred: TensorDict) -> torch.Tensor:
     # energy: [n_graphs, ]
     configs_weight = ref.weight.view(-1, 1, 1)  # [n_graphs, ]
     configs_stress_weight = ref.stress_weight.view(-1, 1, 1)  # [n_graphs, ]
-    num_atoms = (ref.ptr[1:] - ref.ptr[:-1]).view(-1, 1, 1)  # [n_graphs,]
     return torch.mean(
         configs_weight
         * configs_stress_weight
-        * torch.square((ref["stress"] - pred["stress"]) / num_atoms)
+        * torch.square(ref["stress"] - pred["stress"])
     )  # []