Fix

lightly-ai · Nov 20, 2023 · 4c57c10 · 4c57c10
1 parent c916d1e
commit 4c57c10
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/lightly/models/modules/heads.py b/lightly/models/modules/heads.py
@@ -371,9 +371,11 @@ 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(
+            ] + (1 - self.beta) * x[
+                mask
+            ].mean(  # type: ignore
                 axis=0
-            )  # type: ignore
+            )
 
         return group_features
 

diff --git a/lightly/models/modules/nn_memory_bank.py b/lightly/models/modules/nn_memory_bank.py
@@ -10,7 +10,7 @@
 from lightly.loss.memory_bank import MemoryBankModule
 
 
-class NNMemoryBankModule(MemoryBankModule):  # type: ignore # Cannot subclass type Any.
+class NNMemoryBankModule(MemoryBankModule):
     """Nearest Neighbour Memory Bank implementation
 
     This class implements a nearest neighbour memory bank as described in the