move first residual after initial temporal attention

lucidrains · Jul 24, 2022 · 601acee · 601acee
1 parent e4d9e9d
commit 601acee
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'video-diffusion-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.5.3',
+  version = '0.6.0',
   license='MIT',
   description = 'Video Diffusion - Pytorch',
   long_description_content_type = 'text/markdown',

diff --git a/video_diffusion_pytorch/video_diffusion_pytorch.py b/video_diffusion_pytorch/video_diffusion_pytorch.py
@@ -485,10 +485,11 @@ def forward(
         time_rel_pos_bias = self.time_rel_pos_bias(x.shape[2], device = x.device)
 
         x = self.init_conv(x)
-        r = x.clone()
 
         x = self.init_temporal_attn(x, pos_bias = time_rel_pos_bias)
 
+        r = x.clone()
+
         t = self.time_mlp(time) if exists(self.time_mlp) else None
 
         # classifier free guidance