Expected both inputs to be on same device

binary_op(): expected both inputs to be on same device, but input a is on cuda:1 and input b is on cuda:0

class GaussianNoise(nn.Module):
    def __init__(self, sigma=1.0):
        super().__init__()
        self.sigma = sigma
        self.noise = torch.tensor(0.0, device='cuda:0')

    def forward(self, x):
        sampled_noise = self.noise.repeat(*x.size()).normal_(mean=0, std=self.sigma)
        x = x + sampled_noise
        return x

def network():
    layers = [resnet.layer1, GaussianNoise(1.0)]
   return nn.Sequential(*layers)

I'm using multi-gpus.
How to match gpu devices between x and sampled_noise?

Inside the forward pass, you could create sampled_noise from the attributes of x:

sampled_noise = torch.empty_like(x),normal_(mean=0, std=self.sigma)

or push the noise tensor to the device of x:

sampled_noise = ...
sampled_noise = sampled_noise.to(x.device)
1 Like