Merge pull request #314 from dev-jinwoohong/develop

Fix scale calculation in attention mechanism
lucidrains · May 17, 2024 · 356d7ee · 356d7ee
2 parents 9c20670 + e875f37
commit 356d7ee
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/denoising_diffusion_pytorch/attend.py b/denoising_diffusion_pytorch/attend.py
@@ -72,7 +72,7 @@ def flash_attn(self, q, k, v):
 
         if exists(self.scale):
             default_scale = q.shape[-1]
-            q = q * (scale / default_scale)
+            q = q * (self.scale / default_scale)
 
         q, k, v = map(lambda t: t.contiguous(), (q, k, v))