Cast neighbor pairs to long (expected by torch geometric)

torchmd · May 11, 2023 · 47d256d · 47d256d
1 parent f15d6cc
commit 47d256d
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/torchmdnet/models/utils.py b/torchmdnet/models/utils.py
@@ -215,6 +215,7 @@ def forward(
             neighbors = neighbors[:, mask]
             distances = distances[mask]
             distance_vecs = distance_vecs[mask, :]
+        neighbors = neighbors.to(torch.long)
         if self.return_vecs:
             return neighbors, distances, distance_vecs
         else: