RuntimeError when training

[matheushent]

I just run the code for training and I got the traceback:

Traceback (most recent call last):
  File "train.py", line 64, in <module>
    model.optimize_parameters()
  File "C:\Users\mathe\projects\SwapNet\models\warp_model.py", line 177, in optimize_parameters
    super().optimize_parameters()
  File "C:\Users\mathe\projects\SwapNet\models\base_gan.py", line 198, in optimize_parameters
    self.backward_D()
  File "C:\Users\mathe\projects\SwapNet\models\warp_model.py", line 117, in backward_D
    self.loss_D_fake = self.criterion_GAN(pred_fake, False)
  File "C:\Users\mathe\projects\SwapNet\modules\loss.py", line 119, in __call__
    target_tensor = self.get_target_tensor(prediction, target_is_real)
  File "C:\Users\mathe\projects\SwapNet\modules\loss.py", line 101, in get_target_tensor
    target_tensor = GANLoss.rand_between(low, high).to(
  File "C:\Users\mathe\projects\SwapNet\modules\loss.py", line 75, in rand_between
    return rand_func(1) * (high - low) + low
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

It seems some parameters are not allocated on GPU. Any ideia of solving it?

[matheushent]

Solved changing line 75 from modules/loss.py from return rand_func(1) * (high - low) + low to return rand_func(1, device=torch.device('cuda:0')) * (high - low) + low

[matheushent]

This problem would be better solved changing the whole function from

def rand_between(low, high, normal=False):
        """
        Args:
            low: a torch.Tensor
            high: a torch.Tensor
            normal: whether to use normal distribution. if not, will use uniform

        Returns: random tensor between low and high
        """
        rand_func = torch.randn if normal else torch.rand
        return rand_func(1) * (high - low) + low

to

def rand_between(low, high, normal=False):
        """
        Args:
            low: a torch.Tensor
            high: a torch.Tensor
            normal: whether to use normal distribution. if not, will use uniform

        Returns: random tensor between low and high
        """
        if normal:
               return torch.randn(1) * (high - low) + low
        else:
               return torch.rand(1) * (high - low) + low

It seems that rand_func variable is allocated on CPU even the default device is the GPU.