Update blocks.py

ACEsuit · Oct 1, 2024 · ee4f1d9 · ee4f1d9
1 parent fa86566
commit ee4f1d9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mace/modules/blocks.py b/mace/modules/blocks.py
@@ -568,7 +568,7 @@ def _setup(self) -> None:
         # Normalization
         self.normalize = nn.BatchNorm(
             self.irreps_out,
-            affine=False,
+            affine=True,
             reduce="mean",
             instance=True,
             normalization="component",
@@ -650,7 +650,7 @@ def _setup(self) -> None:
         # Normalization
         self.normalize = nn.BatchNorm(
             self.irreps_out,
-            affine=False,
+            affine=True,
             reduce="mean",
             instance=True,
             normalization="component",