quick fix for training on images

lucidrains · Oct 19, 2023 · 433f584 · 433f584
1 parent 49dd0e9
commit 433f584
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/magvit2_pytorch/magvit2_pytorch.py b/magvit2_pytorch/magvit2_pytorch.py
@@ -784,7 +784,7 @@ def forward(
         # accept images for image pretraining (curriculum learning from images to video)
 
         if video_or_images.ndim == 4:
-            video = rearrange(video, 'b c ... -> b c 1 ...')
+            video = rearrange(video_or_images, 'b c ... -> b c 1 ...')
         else:
             video = video_or_images
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'magvit2-pytorch',
   packages = find_packages(),
-  version = '0.0.16',
+  version = '0.0.17',
   license='MIT',
   description = 'MagViT2 - Pytorch',
   long_description_content_type = 'text/markdown',