Tensor device error when computing kl_div in VAE

Hello, @lucidrains 😃

I saw https://github.com/lucidrains/DALLE-pytorch/commit/706f06db0c8cede9fe02d16cbfe8ee09144858ea this commit, and I got below error when I trained with it.

Traceback (most recent call last):
  File "/root/.pycharm_helpers/pydev/pydevd.py", line 1434, in _exec
    pydev_imports.execfile(file, globals, locals)  # execute the script
  File "/root/.pycharm_helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/home/shared/workspace/torch_research/text-to-image/dalle-pytorch/train_VAE.py", line 237, in <module>
    main()
  File "/home/shared/workspace/torch_research/text-to-image/dalle-pytorch/train_VAE.py", line 164, in main
    loss = get_loss(images)
  File "/home/shared/workspace/torch_research/text-to-image/dalle-pytorch/train_VAE.py", line 145, in get_loss
    loss = F.smooth_l1_loss(images, recons) + vae(images, return_loss=True)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/shared/workspace/torch_research/text-to-image/dalle-pytorch/models/model_arch.py", line 173, in forward
    kl_div = (qy * (log_qy - g)).sum(dim = (1, 2)).mean()
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Reason is the device of g construction is not adaptable to the current device.

https://github.com/lucidrains/DALLE-pytorch/blob/706f06db0c8cede9fe02d16cbfe8ee09144858ea/dalle_pytorch/dalle_pytorch.py#L172

This can be solved by

device = img.device
....
g = torch.log(torch.tensor([1. / num_tokens], device=device))

Please check if it is a correct way to fix 😃

[lucidrains]

@johnmccain @kobiso i did correct it in the latest version! https://github.com/lucidrains/DALLE-pytorch/blob/main/dalle_pytorch/dalle_pytorch.py#L172 let me push out another patch just in case

[kobiso]

Thanks for quick revision!