make sure training on only images work

lucidrains · Dec 9, 2022 · 3f4071d · 3f4071d
1 parent 0af63d7
commit 3f4071d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/make_a_video_pytorch/make_a_video.py b/make_a_video_pytorch/make_a_video.py
@@ -289,7 +289,7 @@ def forward(
             x, = unpack(x, ps, '* c h w')
             x = rearrange(x, 'b f c h w -> b c f h w')
 
-        if not exists(self.down_time):
+        if not is_video or not exists(self.down_time):
             return x
 
         x = self.down_time(x)
@@ -330,7 +330,7 @@ def forward(
             x, = unpack(x, ps, '* c h w')
             x = rearrange(x, 'b f c h w -> b c f h w')
 
-        if not exists(self.up_time):
+        if not is_video or not exists(self.up_time):
             return x
 
         x = self.up_time(x)