RuntimeError when running on GPU

[alekitto]

Trying to run this on NVIDIA GPU, it crashes almost immediately with the following error:

$ python main.py --config config/base.yaml --target ~/img.png --experiment experiment_5x1 --signature shield --log_dir log/
Input image shape is: (1024, 1024, 3)
=> Adding [1] paths, [circle] ...
  0%|                                                                                                                                                                             | 0/500 [00:00<?, ?it/s, loss=0]
Traceback (most recent call last):
  File "main.py", line 605, in <module>
    loss.backward()
  File "/home/user/miniconda3/envs/live/lib/python3.7/site-packages/torch/tensor.py", line 166, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/home/user/miniconda3/envs/live/lib/python3.7/site-packages/torch/autograd/__init__.py", line 99, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Function AddBackward0 returned an invalid gradient at index 1 - expected type torch.FloatTensor but got torch.cuda.FloatTensor