fix: reshaping cell tensor for ase calculation if needed

IntelLabs · Nov 18, 2024 · 61bbf1d · 61bbf1d
1 parent 34f5cb8
commit 61bbf1d
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/matsciml/datasets/utils.py b/matsciml/datasets/utils.py
@@ -882,7 +882,9 @@ def _all_sites_have_neighbors(neighbors):
         "pos": coords,
     }
 
-    cell = rearrange(cell, "i j -> () i j")
+    # only do the reshape if we are missing a dimension
+    if cell.ndim == 2:
+        cell = rearrange(cell, "i j -> () i j")
     return_dict["offsets"] = einsum(return_dict["images"], cell, "v i, n i j -> v j")
     src, dst = return_dict["src_nodes"], return_dict["dst_nodes"]
     return_dict["unit_offsets"] = (