Fix max_num_pairs incorrect error checking with negative max_num_pairs

torchmd · May 11, 2023 · f15d6cc · f15d6cc
1 parent d3578f9
commit f15d6cc
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/torchmdnet/models/utils.py b/torchmdnet/models/utils.py
@@ -96,7 +96,7 @@ def __init__(
         self,
         cutoff_lower=0.0,
         cutoff_upper=5.0,
-        max_num_pairs=32,
+        max_num_pairs=-32,
         return_vecs=False,
         loop=False,
         strategy="brute",
@@ -203,10 +203,10 @@ def forward(
             use_periodic=self.use_periodic,
         )
         if self.check_errors:
-            if num_pairs[0] > self.max_num_pairs:
+            if num_pairs[0] > max_pairs:
                 raise RuntimeError(
                     "Found num_pairs({}) > max_num_pairs({})".format(
-                        num_pairs[0], self.max_num_pairs
+                        num_pairs[0], max_pairs
                     )
                 )
         # Remove (-1,-1)  pairs