learning/annotated_deep_learning_paper_implementations
1
0
Fork 0
mirror of https://github.com/labmlai/annotated_deep_learning_paper_implementations.git synced 2025-10-29 09:38:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
Varuna Jayasiri
2022-09-24 14:39:10 +05:30
parent de36f9b6be
commit eb92824e58
2 changed files with 59 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
labml_nn/diffusion/stable_diffusion/model/unet_attention.py
View File

@ -173,6 +173,8 @@ class CrossAttention(nn.Module):
k = self.to_k(cond)
v = self.to_v(cond)
print('use flash', CrossAttention.use_flash_attention)
if CrossAttention.use_flash_attention and self.flash is not None and cond is None and self.d_head <= 128:
return self.flash_attention(q, k, v)
else:
@ -186,7 +188,7 @@ class CrossAttention(nn.Module):
"""
print('flash')
# Get batch size and number of elements along sequence axis (width * height)
batch_size, seq_len, _ = q.shape