Reformatting MACE (ORNL#307)

* reformatting MACE * move irreps_string tool * Comment * adjusting position accounting for batch
RylieWeaver · Nov 12, 2024 · ab04752 · ab04752
1 parent 8a66236
commit ab04752
Show file tree

Hide file tree

Showing 4 changed files with 614 additions and 457 deletions.
diff --git a/hydragnn/models/Base.py b/hydragnn/models/Base.py
@@ -20,6 +20,7 @@
 import sys
 from hydragnn.utils.distributed import get_device
 from hydragnn.utils.print.print_utils import print_master
+from hydragnn.utils.model.operations import get_edge_vectors_and_lengths
 
 import inspect
 
@@ -136,6 +137,10 @@ def _init_conv(self):
             self.feature_layers.append(BatchNorm(self.hidden_dim))
 
     def _embedding(self, data):
+        if not hasattr(data, "edge_shifts"):
+            data.edge_shifts = torch.zeros(
+                (data.edge_index.size(1), 3), device=data.edge_index.device
+            )
         conv_args = {"edge_index": data.edge_index.to(torch.long)}
         if self.use_edge_attr:
             assert (