quick fix

lucidrains · Mar 19, 2023 · bde9f22 · bde9f22
1 parent 5e28401
commit bde9f22
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/make_a_video_pytorch/make_a_video.py b/make_a_video_pytorch/make_a_video.py
@@ -277,6 +277,7 @@ def __init__(
         self.temporal_attn = Attention(dim = dim, dim_head = dim_head, heads = heads)
         self.temporal_rel_pos_bias = ContinuousPositionBias(dim = dim // 2, heads = heads, num_dims = 1)
 
+        self.has_feed_forward = add_feed_forward
         if not add_feed_forward:
             return
 
@@ -315,7 +316,9 @@ def forward(
 
             x = rearrange(x, '(b h w) f c -> b c f h w', w = w, h = h)
 
-        x = self.ff(x, enable_time = enable_time) + x
+        if self.has_feed_forward:
+            x = self.ff(x, enable_time = enable_time) + x
+
         return x
 
 # resnet block

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'make-a-video-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.1.0',
+  version = '0.1.1',
   license='MIT',
   description = 'Make-A-Video - Pytorch',
   author = 'Phil Wang',