Fix

lightly/models/modules/heads.py

@@ -371,11 +371,9 @@ def get_updated_group_features(self, x: torch.Tensor) -> torch.Tensor:
             mask = assignments == assigned_class
             group_features[assigned_class] = self.beta * self.group_features[
                 assigned_class
-            ] + (1 - self.beta) * x[
-                mask
-            ].mean(  # type: ignore
+            ] + (1 - self.beta) * x[mask].mean(
                 axis=0
-            )
+            )  # type: ignore
 
         return group_features
 

lightly/models/modules/nn_memory_bank.py

@@ -21,8 +21,7 @@ class NNMemoryBankModule(MemoryBankModule):
 
     Attributes:
         size:
-            Number of keys the memory bank can store. If set to 0,
-            memory bank is not used.
+            Number of keys the memory bank can store.
 
     Examples:
         >>> model = NNCLR(backbone)
@@ -40,6 +39,8 @@ class NNMemoryBankModule(MemoryBankModule):
     """
 
     def __init__(self, size: int = 2**16):
+        if size <= 0:
+            raise ValueError(f"Memory bank size must be positive, got {size}.")
         super(NNMemoryBankModule, self).__init__(size)
 
     def forward(
@@ -57,6 +58,7 @@ def forward(
         """
 
         output, bank = super(NNMemoryBankModule, self).forward(output, update=update)
+        assert bank is not None
         bank = bank.to(output.device).t()
 
         output_normed = torch.nn.functional.normalize(output, dim=1)