Final changes

Author: frostedoyster

src/metatrain/experimental/phace/model.py

@@ -138,9 +138,9 @@ def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo) -> None:
                 cgs,
                 irreps_in=equivariant_message_passer.irreps_out,
                 # TODO: speed up with something like this?
-                # requested_LS_string=(
-                #     "0_1" if idx == self.num_message_passing_layers - 2 else None
-                # ),
+                requested_LS_string=(
+                    "0_1" if idx == self.num_message_passing_layers - 2 else None
+                ),
             )
             generalized_cg_iterators.append(generalized_cg_iterator)
         self.equivariant_message_passers = torch.nn.ModuleList(
@@ -499,7 +499,7 @@ def _add_output(self, target_name: str, target_info: TargetInfo) -> None:
                 # TODO: the equivariant could be a scalar...
         else:
             # specified by the user
-            use_mlp = self.head_types[target_name] == "mlp"
+            use_mlp = (self.head_types[target_name] == "mlp")
 
         self.outputs[target_name] = ModelOutput(
             quantity=target_info.quantity,

src/metatrain/experimental/phace/modules/layers.py

@@ -23,11 +23,8 @@ def __init__(self, n_inputs: int) -> None:
         self.mlp = torch.nn.Sequential(
             Linear(n_inputs, 4 * n_inputs),
             torch.nn.SiLU(),
-            # Linear(4*n_inputs, 4*n_inputs),
-            # torch.nn.SiLU(),
-            # Linear(4*n_inputs, 4*n_inputs),
-            # torch.nn.SiLU(),
             Linear(4 * n_inputs, n_inputs),
+            torch.nn.SiLU(),
         )
 
     def forward(self, features: TensorMap) -> TensorMap: