asserts

lucidrains · Oct 15, 2023 · 4768d18 · 4768d18
1 parent 89a29dd
commit 4768d18
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/magvit2_pytorch/magvit2_pytorch.py b/magvit2_pytorch/magvit2_pytorch.py
@@ -446,12 +446,13 @@ def forward(
         return_loss = False,
         return_codes = False
     ):
-        is_images = video_or_images.ndim == 4
+        assert not (return_loss and return_codes)
+        assert video_or_images.ndim in {4, 5}
 
         # accept images for image pretraining (curriculum learning from images to video)
 
-        if is_images:
-            video = rearrange(video_or_images, 'b c ... -> b c 1 ...')
+        if video_or_images.ndim == 4:
+            video = rearrange(video, '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.2',
+  version = '0.0.3',
   license='MIT',
   description = 'MagViT2 - Pytorch',
   long_description_content_type = 'text/markdown',