fixed prototype norm in swav (#249)

vturrisi · Jun 4, 2022 · 16183cf · 16183cf
1 parent 4d86dc9
commit 16183cf
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/solo/methods/swav.py b/solo/methods/swav.py
@@ -78,6 +78,8 @@ def __init__(
         self.prototypes = nn.utils.weight_norm(
             nn.Linear(proj_output_dim, num_prototypes, bias=False)
         )
+        self.prototypes.weight_g.data.fill_(1)  # type: ignore
+        self.prototypes.weight_g.requires_grad = False
 
     @staticmethod
     def add_model_specific_args(parent_parser: argparse.ArgumentParser) -> argparse.ArgumentParser: