Apply fix

lab-cosmo · Apr 4, 2024 · 18c40b5 · 18c40b5
1 parent bf44a2b
commit 18c40b5
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/src/metatensor/models/experimental/soap_bpnn/model.py b/src/metatensor/models/experimental/soap_bpnn/model.py
@@ -18,6 +18,14 @@
 DEFAULT_MODEL_HYPERS = DEFAULT_HYPERS["model"]
 
 
+class Identity(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: TensorMap) -> TensorMap:
+        return x
+
+
 class MLPMap(torch.nn.Module):
     def __init__(self, all_species: List[int], hypers: dict) -> None:
         super().__init__()
@@ -248,7 +256,7 @@ def __init__(
         if hypers_bpnn["layernorm"]:
             self.layernorm = LayerNormMap(self.all_species, soap_size)
         else:
-            self.layernorm = torch.nn.Identity()
+            self.layernorm = Identity()
 
         self.bpnn = MLPMap(self.all_species, hypers_bpnn)
 

diff --git a/src/metatensor/models/experimental/soap_bpnn/tests/test_torchscript.py b/src/metatensor/models/experimental/soap_bpnn/tests/test_torchscript.py
@@ -47,7 +47,7 @@ def test_torchscript_with_identity():
         },
     )
     hypers = copy.deepcopy(DEFAULT_HYPERS["model"])
-    hypers["model"]["bpnn"]["layernorm"] = False
+    hypers["bpnn"]["layernorm"] = False
     soap_bpnn = Model(capabilities, hypers)
     soap_bpnn = torch.jit.script(soap_bpnn)