getting right variables

vturrisi · Jan 8, 2024 · 140c8be · 140c8be
1 parent 882854a
commit 140c8be
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/solo/methods/all4one.py b/solo/methods/all4one.py
@@ -301,7 +301,7 @@ def training_step(self, batch: Sequence[Any], batch_idx: int) -> torch.Tensor:
         out = super().training_step(batch, batch_idx)
         class_loss = out["loss"]
         feats1, feats2 = out["feats"]
-        momentum_z1, momentum_z2 = out["momentum_feats"]
+        momentum_z1, momentum_z2 = out["momentum_z"]
 
         z1 = self.projector(feats1)
         z2 = self.projector(feats2)