diff --git a/src/metatrain/experimental/phace/modules/cg.py b/src/metatrain/experimental/phace/modules/cg.py
index 8d5d5c78..edc4b98e 100644
--- a/src/metatrain/experimental/phace/modules/cg.py
+++ b/src/metatrain/experimental/phace/modules/cg.py
@@ -88,7 +88,7 @@ def cg_combine_l1l2L(tensor12, cg_tensor):
     )
     return out_tensor.swapaxes(
         1, 2
-    )  # / (cg_tensor.shape[0]*cg_tensor.shape[1]*cg_tensor.shape[2])
+    )
 
 
 def get_cg_coefficients(l_max):
diff --git a/src/metatrain/experimental/phace/tests/test_regression.py b/src/metatrain/experimental/phace/tests/test_regression.py
index 6ccea441..1c1442a6 100644
--- a/src/metatrain/experimental/phace/tests/test_regression.py
+++ b/src/metatrain/experimental/phace/tests/test_regression.py
@@ -30,6 +30,7 @@ def test_regression_init():
         length_unit="Angstrom", atomic_types=[1, 6, 7, 8], targets=targets
     )
     model = PhACE(MODEL_HYPERS, dataset_info)
+    model = torch.jit.script(model)
 
     # Predict on the first five systems
     systems = read_systems(DATASET_PATH)[:5]
@@ -106,6 +107,7 @@ def test_regression_train():
     systems = [system.to(torch.float32) for system in systems]
     for system in systems:
         get_system_with_neighbor_lists(system, model.requested_neighbor_lists())
+    model = torch.jit.script(model)
     output = model(
         systems[:5],
         {"mtt::U0": ModelOutput(quantity="energy", unit="", per_atom=False)},