Fix bug in TN (#199)

A dX**2 operation should be a torch.matmul(dX,dX).
torchmd · Jul 3, 2023 · 62d2885 · 62d2885
1 parent a116847
commit 62d2885
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchmdnet/models/tensornet.py b/torchmdnet/models/tensornet.py
@@ -402,7 +402,7 @@ def forward(self, X, edge_index, edge_weight, edge_attr):
         A = self.linears_tensor[4](A.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
         S = self.linears_tensor[5](S.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
         dX = I + A + S
-        X = X + dX + dX**2
+        X = X + dX + torch.matmul(dX,dX)
         return X
 
     def message(self, I_j, A_j, S_j, edge_attr):