steffen/ml_lib
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

CCS intergration dataloader

Browse Source
This commit is contained in:
Steffen
2021-03-19 17:17:16 +01:00
parent 8e719af554
commit 43cf0ad00d
1 changed files with 3 additions and 4 deletions

5
modules/blocks.py

@ -263,9 +263,8 @@ class Attention(nn.Module):
assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions' assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions'
mask = mask[:, None, :] * mask[:, :, None] mask = mask[:, None, :] * mask[:, :, None]
mask = repeat(mask, 'b n d -> b h n d', h=h) # My addition mask = repeat(mask, 'b n d -> b h n d', h=h) # My addition
dots.masked_fill_(~mask, mask_value)
#dots.masked_fill_(~mask, mask_value) # dots.masked_fill_(mask, mask_value) # My addition
dots.masked_fill_(mask, mask_value)
del mask del mask
attn = dots.softmax(dim=-1) attn = dots.softmax(dim=-1)