address #44

lucidrains · Jul 23, 2024 · 06edbd3 · 06edbd3
1 parent e085ec3
commit 06edbd3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/magvit2_pytorch/magvit2_pytorch.py b/magvit2_pytorch/magvit2_pytorch.py
@@ -1233,12 +1233,12 @@ def __init__(
 
                 encoder_layer = Sequential(
                     Residual(TokenShift(TimeAttention(**attn_kwargs))),
-                    Residual(TokenShift(FeedForward(dim, dim_cond = dim_cond)))
+                    Residual(TokenShift(FeedForward(dim)))
                 )
 
                 decoder_layer = Sequential(
                     Residual(TokenShift(TimeAttention(**attn_kwargs))),
-                    Residual(TokenShift(FeedForward(dim, dim_cond = dim_cond)))
+                    Residual(TokenShift(FeedForward(dim)))
                 )
 
             elif layer_type == 'cond_attend_space':
@@ -1255,12 +1255,12 @@ def __init__(
 
                 encoder_layer = Sequential(
                     Residual(SpaceAttention(**attn_kwargs)),
-                    Residual(FeedForward(dim))
+                    Residual(FeedForward(dim, dim_cond = dim_cond))
                 )
 
                 decoder_layer = Sequential(
                     Residual(SpaceAttention(**attn_kwargs)),
-                    Residual(FeedForward(dim))
+                    Residual(FeedForward(dim, dim_cond = dim_cond))
                 )
 
             elif layer_type == 'cond_linear_attend_space':

diff --git a/magvit2_pytorch/version.py b/magvit2_pytorch/version.py
@@ -1 +1 @@
-__version__ = '0.4.8'
+__version__ = '0.4.9'