rename a keyword

Author: lucidrains

README.md

@@ -43,7 +43,7 @@ vae = DiscreteVAE(
 
 images = torch.randn(4, 3, 256, 256)
 
-loss = vae(images, return_recon_loss = True)
+loss = vae(images, return_loss = True)
 loss.backward()
 
 # train with a lot of data to learn a good codebook

dalle_pytorch/dalle_pytorch.py

@@ -143,7 +143,7 @@ def decode(
     def forward(
         self,
         img,
-        return_recon_loss = False,
+        return_loss = False,
         return_logits = False
     ):
         num_tokens, kl_div_loss_weight = self.num_tokens, self.kl_div_loss_weight
@@ -157,7 +157,7 @@ def forward(
         sampled = einsum('b n h w, n d -> b d h w', soft_one_hot, self.codebook.weight)
         out = self.decoder(sampled)
 
-        if not return_recon_loss:
+        if not return_loss:
             return out
 
         # reconstruction loss

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'dalle-pytorch',
   packages = find_packages(),
-  version = '0.1.0',
+  version = '0.1.1',
   license='MIT',
   description = 'DALL-E - Pytorch',
   author = 'Phil Wang',