Hi, I am building a kind of Generator Network in this I am getting the same error. I am new to PyTorch Framework
Error is shown:
/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py in batch_norm(input, running_mean, running_var, weight, bias, training, momentum, eps)
   2056     return torch.batch_norm(
   2057         input, weight, bias, running_mean, running_var,
-> 2058         training, momentum, eps, torch.backends.cudnn.enabled
   2059     )
   2060 

RuntimeError: Tensor for argument #2 'weight' is on CPU, but expected it to be on GPU (while checking arguments for cudnn_batch_norm)

class netG(nn.Module):
  def __init__(self,channel_rate=64,drop_rate=0.0):
    super(netG, self).__init__()
    self.channel_rate = channel_rate
    self.nChannels = 4*self.channel_rate
    self.drop_rate = drop_rate
    self.dilation_factor = 1
    self.head = nn.Conv2d(3,4*self.channel_rate,3,1,1)
  
  def dense_block(self,x):
    block = nn.Sequential(
    nn.BatchNorm2d(self.nChannels),
    nn.LeakyReLU(0.2), 
    nn.Conv2d(self.nChannels,4*self.channel_rate,1,1),
    nn.BatchNorm2d(4*self.channel_rate),
    nn.Conv2d(4*self.channel_rate,self.channel_rate,3,1,self.dilation_factor,self.dilation_factor),
    nn.BatchNorm2d(self.channel_rate),
    nn.Dropout2d(self.drop_rate)
    )
    return block(x)


  def tail(self,x):
    block = nn.Sequential(
    nn.LeakyReLU(0.2),
    nn.Conv2d(self.nChannels,4*self.channel_rate,1,1),
    nn.BatchNorm2d(4*self.channel_rate),
    nn.Dropout(self.drop_rate)
    )
    return block(x)
    
  def last_layer(self,x):
    block= nn.Sequential(
    nn.Conv2d(self.nChannels,self.channel_rate,3,1,1),
    nn.PReLU(),
    nn.Conv2d(self.channel_rate,3,3,1,1),
    nn.Tanh()
    )
    return block(x)

  def forward(self,x):
    x = self.head(x)
    x1 = x
    self.nChannels = x.size(1)
    self.dilation_factor = 1
    d1 = self.dense_block(x)
    x = torch.cat((x,d1),1)
    self.dilation_factor = 1
    self.nChannels = x.size(1)
    d2 = self.dense_block(x)
    x = torch.cat((x,d2),1)
    self.dilation_factor = 2
    self.nChannels  = x.size(1)
    d4 = self.dense_block(x)
    x = torch.cat((x,d4),1)
    self.dilation_factor = 1
    self.nChannels  = x.size(1)
    d5 = self.dense_block(x)
    x = torch.cat((x,d5),1)
    self.dilation_factor = 3
    self.nChannels  = x.size(1)
    d6 = self.dense_block(x)
    x = torch.cat((x,d6),1)
    self.dilation_factor = 1
    self.nChannels  = x.size(1)
    d7 = self.dense_block(x)
    x = torch.cat((x,d7),1)
    self.dilation_factor = 2
    self.nChannels  = x.size(1)
    d8 = self.dense_block(x)
    x = torch.cat((x,d8),1)
    self.dilation_factor = 1
    self.nChannels  = x.size(1)
    d9 = self.dense_block(x)
    x = torch.cat((x,d9),1)
    self.dilation_factor = 1
    self.nChannels  = x.size(1)
    d10 = self.dense_block(x)
    x = self.tail(x)
    x = torch.cat((x,x1),1)
    self.nChannels  = x.size(1)
    x = self.last_layer(x)
    return x