Fix incorrect type

torchmd · Mar 20, 2024 · 500a8e4 · 500a8e4
1 parent eeded26
commit 500a8e4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchmdnet/extensions/__init__.py b/torchmdnet/extensions/__init__.py
@@ -137,10 +137,10 @@ def get_neighbor_pairs_fwd_meta(
 ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
     """Returns empty vectors with the correct shape for the output of get_neighbor_pairs_kernel."""
     size = max_num_pairs
-    edge_index = torch.empty((2, size), dtype=torch.long, device=positions.device)
+    edge_index = torch.empty((2, size), dtype=torch.int, device=positions.device)
     edge_distance = torch.empty((size,), dtype=positions.dtype, device=positions.device)
     edge_vec = torch.empty((size, 3), dtype=positions.dtype, device=positions.device)
-    num_pairs = torch.empty((1,), dtype=torch.long, device=positions.device)
+    num_pairs = torch.empty((1,), dtype=torch.int, device=positions.device)
     return edge_index, edge_vec, edge_distance, num_pairs