[Fix]: Fix UT

open-mmlab · Apr 27, 2022 · cd7f798 · cd7f798
1 parent d26bfef
commit cd7f798
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/mmselfsup/models/algorithms/cae.py b/mmselfsup/models/algorithms/cae.py
@@ -77,7 +77,8 @@ def forward_train(self, samples: Sequence, **kwargs) -> dict:
         img, img_target, mask = samples
 
         # normalize images and the images to get the target
-        img = self.img_norm(img)
+        img_list = [self.img_norm(x).unsqueeze(0) for x in img]
+        img = torch.cat(img_list)
         img_target = 0.8 * img_target + 0.1
 
         mask = mask.flatten(1).to(torch.bool)