Fix

lightly-ai · Sep 11, 2023 · 222e12e · 222e12e
1 parent 1cd51eb
commit 222e12e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/lightly/models/modules/heads.py b/lightly/models/modules/heads.py
@@ -717,7 +717,7 @@ def __init__(
             (hidden_dim, hidden_dim, nn.BatchNorm1d(hidden_dim), nn.ReLU())
             for _ in range(num_layers - 2) # Exclude first and last layer.
         ]
-        super(BarlowTwinsProjectionHead, self).__init__(
+        super(VICRegProjectionHead, self).__init__(
             [
                 (input_dim, hidden_dim, nn.BatchNorm1d(hidden_dim), nn.ReLU()),
                 *hidden_layers,