Expected object of type torch.cuda.FloatTensor but found type torch.FloatTensor for argument #3 'other'

[doantientai]

I am trying to run UNIT on my own data but receiving the error below. It seems related to the vgg model. Please help..

Traceback (most recent call last):
  File "train.py", line 65, in <module>
    trainer.gen_update(images_a, images_b, config)
  File "/media/tai/6TB/Projects/UNIT/trainer.py", line 295, in gen_update
    self.loss_gen_vgg_a = self.compute_vgg_loss(self.vgg, x_ba, x_b) if hyperparameters['vgg_w'] > 0 else 0
  File "/media/tai/6TB/Projects/UNIT/trainer.py", line 314, in compute_vgg_loss
    img_vgg = vgg_preprocess(img)
  File "/media/tai/6TB/Projects/UNIT/utils.py", line 244, in vgg_preprocess
    batch = batch.sub(Variable(mean)) # subtract mean
RuntimeError: Expected object of type torch.cuda.FloatTensor but found type torch.FloatTensor for argument #3 'other'

[doantientai]

Sorry I did not try to spam.. I thought my post was not submitted...