Fix DINOProjectionHead BatchNorm Handling

lightly-ai · Nov 8, 2024 · 1c560e0 · 1c560e0
1 parent 94050b3
commit 1c560e0
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/lightly/models/modules/heads.py b/lightly/models/modules/heads.py
@@ -673,12 +673,20 @@ def __init__(
         norm_last_layer: bool = True,
     ):
         """Initializes the DINOProjectionHead with the specified dimensions."""
-        bn = nn.BatchNorm1d(hidden_dim) if batch_norm else None
-
         super().__init__(
             [
-                (input_dim, hidden_dim, bn, nn.GELU()),
-                (hidden_dim, hidden_dim, bn, nn.GELU()),
+                (
+                    input_dim,
+                    hidden_dim,
+                    nn.BatchNorm1d(hidden_dim) if batch_norm else None,
+                    nn.GELU(),
+                ),
+                (
+                    hidden_dim,
+                    hidden_dim,
+                    nn.BatchNorm1d(hidden_dim) if batch_norm else None,
+                    nn.GELU(),
+                ),
                 (hidden_dim, bottleneck_dim, None, None),
             ]
         )