fix: missing max_neighbors in shortcut periodic calculation

IntelLabs · Nov 18, 2024 · 1d219c3 · 1d219c3
1 parent 64e225d
commit 1d219c3
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/matsciml/datasets/transforms/pbc.py b/matsciml/datasets/transforms/pbc.py
@@ -97,7 +97,10 @@ def __call__(self, data: DataDict) -> DataDict:
             structure = data["structure"]
             if isinstance(structure, Structure):
                 graph_props = calculate_periodic_shifts(
-                    structure, self.cutoff_radius, self.adaptive_cutoff
+                    structure,
+                    self.cutoff_radius,
+                    self.adaptive_cutoff,
+                    max_neighbors=self.max_neighbors,
                 )
                 data.update(graph_props)
                 return data