fix the EnergyRefRemover which was accidentally removing the total su…

…m of all atomref energies of the whole batch from the energies
torchmd · Feb 16, 2024 · 500b04c · 500b04c
1 parent 94cc325
commit 500b04c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchmdnet/module.py b/torchmdnet/module.py
@@ -42,9 +42,9 @@ def __init__(self, atomref):
         self._atomref = atomref
 
     def forward(self, data):
-        self._atomref = self._atomref.to(data.z.device)
+        self._atomref = self._atomref.to(data.z.device).type(data.y.dtype)
         if "y" in data:
-            data.y -= self._atomref[data.z].sum()
+            data.y.index_add_(0, data.batch, -self._atomref[data.z])
         return data