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

better naming

Browse Source
This commit is contained in:
Phil Wang
2022-04-14 09:24:31 -07:00
parent f2c52d8239
commit 8b5c2385b0
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
dalle2_pytorch/dalle2_pytorch.py
View File

@@ -320,8 +320,8 @@ class DiffusionPriorNetwork(nn.Module):
# but let's just do it right
if exists(mask):
all_masked_out = mask.any(dim = -1)
mask = torch.cat((mask, rearrange(all_masked_out, 'b -> b 1')), dim = 1)
not_all_masked_out = mask.any(dim = -1)
mask = torch.cat((mask, rearrange(not_all_masked_out, 'b -> b 1')), dim = 1)
if exists(mask):
mask = F.pad(mask, (0, 2), value = True) # extend mask for text embedding, noised image embedding, time step embedding, and learned query