get patch out

lucidrains · Jan 6, 2021 · d40aecc · d40aecc
1 parent 4d4433e
commit d40aecc
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/dalle_pytorch/dalle_pytorch.py b/dalle_pytorch/dalle_pytorch.py
@@ -215,9 +215,11 @@ def __init__(
         self.image_seq_len = image_seq_len
         self.total_tokens = num_text_tokens + num_image_tokens + 1 # extra for EOS
 
-        if self.vae is not None:
+        self.vae = vae
+        if exists(self.vae):
             self.vae = vae
             self.image_emb = vae.codebook
+
         self.transformer = Decoder(dim = dim, depth = depth, heads = heads)
 
         self.to_logits = nn.Sequential(

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'dalle-pytorch',
   packages = find_packages(),
-  version = '0.0.4',
+  version = '0.0.5',
   license='MIT',
   description = 'DALL-E - Pytorch',
   author = 'Phil Wang',