rescale values in linear attention to mitigate overflows in fp16 setting

Author: lucidrains

dalle2_pytorch/dalle2_pytorch.py

@@ -1503,6 +1503,7 @@ def forward(self, fmap):
         k = k.softmax(dim = -2)
 
         q = q * self.scale
+        v = v / (x * y)
 
         context = einsum('b n d, b n e -> b d e', k, v)
         out = einsum('b n d, b d e -> b n e', q, context)

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.1.0'
+__version__ = '1.2.0'