fix test setup

pytorch · Oct 7, 2022 · d8945e6 · d8945e6
1 parent ad4d424
commit d8945e6
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/test/prototype_common_utils.py b/test/prototype_common_utils.py
@@ -559,7 +559,7 @@ def make_video_loader(
     num_frames = int(torch.randint(1, 5, ())) if num_frames == "random" else num_frames
 
     def fn(shape, dtype, device):
-        video = make_image(size=shape[-2:], color_space=color_space, extra_dims=shape[:-2], dtype=dtype, device=device)
+        video = make_image(size=shape[-2:], color_space=color_space, extra_dims=shape[:-3], dtype=dtype, device=device)
         return features.Video(video, color_space=color_space)
 
     return VideoLoader(

diff --git a/test/test_prototype_transforms.py b/test/test_prototype_transforms.py
@@ -157,11 +157,8 @@ def test_mixup_cutmix(self, transform, input):
                             features.ColorSpace.RGB,
                         ],
                         dtypes=[torch.uint8],
-                        **(
-                            dict(num_frames=[1, "random"], extra_dims=[()])
-                            if fn is make_videos
-                            else dict(extra_dims=[(4,)])
-                        ),
+                        extra_dims=[(), (4,)],
+                        **(dict(num_frames=["random"]) if fn is make_videos else dict()),
                     )
                     for fn in [
                         make_images,