fixing device, argparser and normalization

ACEsuit · Oct 30, 2024 · ac44b5e · ac44b5e
1 parent 42b9728
commit ac44b5e
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git a/mace/modules/blocks.py b/mace/modules/blocks.py
@@ -694,12 +694,12 @@ def _setup(self) -> None:
             self.reshape = reshape_irreps(self.irreps_out)
 
         elif self.tensor_format in ["non_symmetric_cp", "non_symmetric_tucker"]:
-            self.linear = []
+            self.linear = torch.nn.ModuleList([])
             # Selector TensorProduct
             self.skip_tp = o3.FullyConnectedTensorProduct(
                 self.node_feats_irreps, self.node_attrs_irreps, self.hidden_irreps
             )
-            self.reshape = []
+            self.reshape = torch.nn.ModuleList([])
             for _ in range(self.correlation):
                 self.linear.append(o3.Linear(
                     irreps_mid, self.irreps_out, internal_weights=True, shared_weights=True
@@ -744,7 +744,7 @@ def forward(
                 message = torch.cat((message, _message), dim = -1)
             print("shape of message: ", message.shape)
             return (
-                message, 
+                message / self.avg_num_neighbors, 
                 sc
             )
 

diff --git a/mace/modules/symmetric_contraction.py b/mace/modules/symmetric_contraction.py
@@ -322,7 +322,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor):
                                 +[f"bk{out_channel_idx[nu-1]}"]
                                 +[f"->bk{out_channel_idx[:nu]}"]
                                 ),
-                        *[c_tensor for _ in range(nu)])
+                        *[c_tensor for _ in range(nu)]) / torch.jit._builtins.math.factorial(nu)
 
                 else:
                     if self.tensor_format == "non_symmetric_tucker":
@@ -334,7 +334,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor):
                                 +[f"bk{out_channel_idx[nu-1]}"]
                                 +[f"->bk{out_channel_idx[:nu]}"]
                                 ),
-                        *[c_tensor for _ in range(nu)])
+                        *[c_tensor for _ in range(nu)]) / torch.jit._builtins.math.factorial(nu)
 
 
                 outs[nu] = torch.einsum(
@@ -357,6 +357,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor):
                 # reshape kLM
                 outs[nu] = outs[nu].reshape(outs[nu].shape[0], -1)
 
+            # / factorial(nu) because of extra work done for convenience
             return torch.cat([outs[nu] for nu in range(self.correlation, 0, -1)], dim = 1)
 
         elif "cp" in self.tensor_format:

diff --git a/mace/tools/arg_parser.py b/mace/tools/arg_parser.py
@@ -251,6 +251,7 @@ def build_default_arg_parser() -> argparse.ArgumentParser:
         choices=["symmetric_cp",
                  "non_symmetric_cp",
                  "symmetric_tucker",
+                 "non_symmetric_tucker",
                 ]
     )