refactor: use long instead of int

laserkelvin · Jul 15, 2024 · 8e11437 · 8e11437
1 parent 35e44e9
commit 8e11437
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/matsciml/models/base.py b/matsciml/models/base.py
@@ -346,8 +346,7 @@ def read_batch(self, batch: BatchDict) -> DataDict:
             sizes.append(len(dst_nodes))
             # carve out neighborhoods as dictated by the dataset/transform definition
             sample_pc_pos = (
-                sample[src_nodes.to(int)][None, :]
-                - sample[dst_nodes.to(int).to(int)][:, None]
+                sample[src_nodes.long()][None, :] - sample[dst_nodes.long()][:, None]
             )
             pc_pos.append(sample_pc_pos)
         # pad the position result