fix device error bug

Author: lucidrains

dalle_pytorch/dalle_pytorch.py

@@ -146,7 +146,7 @@ def forward(
         return_loss = False,
         return_logits = False
     ):
-        num_tokens, kl_div_loss_weight = self.num_tokens, self.kl_div_loss_weight
+        device, num_tokens, kl_div_loss_weight = img.device, self.num_tokens, self.kl_div_loss_weight
 
         logits = self.encoder(img)
 
@@ -169,7 +169,7 @@ def forward(
         logits = rearrange(logits, 'b n h w -> b (h w) n')
         qy = F.softmax(logits, dim = -1)
         log_qy = torch.log(qy + 1e-20)
-        g = torch.log(torch.Tensor([1. / num_tokens]))
+        g = torch.log(torch.Tensor([1. / num_tokens], device = device))
         kl_div = (qy * (log_qy - g)).sum(dim = (1, 2)).mean()
 
         return recon_loss + (kl_div * kl_div_loss_weight)

setup.py

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