fix: cross attention skip connection

archinetai · Sep 25, 2022 · 42bf374 · 42bf374
1 parent d60eefa
commit 42bf374
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/audio_diffusion_pytorch/modules.py b/audio_diffusion_pytorch/modules.py
@@ -421,7 +421,7 @@ def __init__(
     def forward(self, x: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
         x = self.attention(x) + x
         if self.use_cross_attention:
-            x = self.cross_attention(x, context=context)
+            x = self.cross_attention(x, context=context) + x
         x = self.feed_forward(x) + x
         return x
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
     name="audio-diffusion-pytorch",
     packages=find_packages(exclude=[]),
-    version="0.0.44",
+    version="0.0.45",
     license="MIT",
     description="Audio Diffusion - PyTorch",
     long_description_content_type="text/markdown",