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

cast attention matrix back to original dtype pre-softmax in attention

Browse Source
This commit is contained in:
Phil Wang
2022-08-20 10:56:01 -07:00
parent 7762edd0ff
commit 083508ff8e
2 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
dalle2_pytorch/dalle2_pytorch.py
View File

@@ -879,6 +879,8 @@ class Attention(nn.Module):
# attention # attention
attn = sim.softmax(dim = -1, dtype = torch.float32) attn = sim.softmax(dim = -1, dtype = torch.float32)
attn = attn.type(sim.dtype)
attn = self.dropout(attn) attn = self.dropout(attn)
# aggregate values # aggregate values
@@ -1637,6 +1639,7 @@ class CrossAttention(nn.Module):
sim = sim.masked_fill(~mask, max_neg_value) sim = sim.masked_fill(~mask, max_neg_value)
attn = sim.softmax(dim = -1, dtype = torch.float32) attn = sim.softmax(dim = -1, dtype = torch.float32)
attn = attn.type(sim.dtype)
out = einsum('b h i j, b h j d -> b h i d', attn, v) out = einsum('b h i j, b h j d -> b h i d', attn, v)
out = rearrange(out, 'b h n d -> b n (h d)') out = rearrange(out, 'b h n d -> b n (h d)')

2
dalle2_pytorch/version.py
View File

@@ -1 +1 @@
__version__ = '1.8.3' __version__ = '1.8.4'