fix bug with random video frame selection for discriminator losses in…

… cvivit, thanks to @1geek0 !
lucidrains · Jan 9, 2023 · 6cbc509 · 6cbc509
1 parent 686937b
commit 6cbc509
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/phenaki_pytorch/cvivit.py b/phenaki_pytorch/cvivit.py
@@ -217,6 +217,7 @@ def pick_video_frame(video, frame_indices):
     batch, device = video.shape[0], video.device
     video = rearrange(video, 'b c f ... -> b f c ...')
     batch_indices = torch.arange(batch, device = device)
+    batch_indices = rearrange(batch_indices, 'b -> b 1')
     images = video[batch_indices, frame_indices]
     images = rearrange(images, 'b 1 c ... -> b c ...')
     return images

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'phenaki-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.66',
+  version = '0.0.67',
   license='MIT',
   description = 'Phenaki - Pytorch',
   author = 'Phil Wang',