Fix

lightly-ai · Nov 20, 2023 · 2fc656b · 2fc656b
1 parent abbb9cd
commit 2fc656b
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/lightly/models/modules/heads.py b/lightly/models/modules/heads.py
@@ -371,9 +371,7 @@ 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(
-                axis=0
-            )  # type: ignore
+            ] + (1 - self.beta) * x[mask].mean(dim=0)
 
         return group_features