wassname/denoising-diffusion-pytorch
1
0
Fork 0
mirror of https://github.com/wassname/denoising-diffusion-pytorch.git synced 2026-09-12 12:22:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

rescale values to prevent linear attention from overflowing in fp16 setting

Browse Source
This commit is contained in:
Phil Wang
2022-07-27 12:25:40 -07:00
parent 6b504c4ae9
commit 12f95b33d8
2 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+2
View File
@@ -211,6 +211,8 @@ class LinearAttention(nn.Module):
k = k.softmax(dim = -1)
q = q * self.scale
v = v / (h * w)
context = torch.einsum('b h d n, b h e n -> b h d e', k, v)
out = torch.einsum('b h d e, b h d n -> b h e n', context, q)