aljaz/DALLE2-pytorch
1
1
Fork 0
mirror of https://github.com/lucidrains/DALLE2-pytorch.git synced 2025-12-19 09:44:19 +01:00
Code Issues Packages Projects Releases Wiki Activity

rescale values in linear attention to mitigate overflows in fp16 setting

Browse Source
This commit is contained in:
Phil Wang
2022-07-27 12:27:32 -07:00
parent 2e35a9967d
commit 07abfcf45b
2 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
dalle2_pytorch/dalle2_pytorch.py
View File

@@ -1503,6 +1503,7 @@ class LinearAttention(nn.Module):
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)