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

fix amp issue for https://github.com/lucidrains/DALLE2-pytorch/issues/82

Browse Source
This commit is contained in:
Phil Wang
2022-05-11 08:21:39 -07:00
parent 2f3c02dba8
commit 3dda2570ed
2 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
dalle2_pytorch/dalle2_pytorch.py
View File

@@ -1492,11 +1492,12 @@ class Unet(nn.Module):
if self.cond_on_image_embeds: if self.cond_on_image_embeds:
image_tokens = self.image_to_cond(image_embed) image_tokens = self.image_to_cond(image_embed)
null_image_embed = self.null_image_embed.to(image_tokens.dtype) # for some reason pytorch AMP not working
image_tokens = torch.where( image_tokens = torch.where(
image_keep_mask, image_keep_mask,
image_tokens, image_tokens,
self.null_image_embed null_image_embed
) )
# take care of text encodings (optional) # take care of text encodings (optional)
@@ -1520,10 +1521,12 @@ class Unet(nn.Module):
text_mask = rearrange(text_mask, 'b n -> b n 1') text_mask = rearrange(text_mask, 'b n -> b n 1')
text_keep_mask = text_mask & text_keep_mask text_keep_mask = text_mask & text_keep_mask
null_text_embed = self.null_text_embed.to(text_tokens.dtype) # for some reason pytorch AMP not working
text_tokens = torch.where( text_tokens = torch.where(
text_keep_mask, text_keep_mask,
text_tokens, text_tokens,
self.null_text_embed null_text_embed
) )
# main conditioning tokens (c) # main conditioning tokens (c)

2
setup.py
View File

@@ -10,7 +10,7 @@ setup(
'dream = dalle2_pytorch.cli:dream' 'dream = dalle2_pytorch.cli:dream'
], ],
}, },
version = '0.2.9', version = '0.2.10',
license='MIT', license='MIT',
description = 'DALL-E 2', description = 'DALL-E 2',
author = 'Phil Wang', author = 'Phil Wang',