Unbale to solve RuntimeError: Expected object of backend CPU but got backend CUDA for argument #2 'weight'

Greetings pytoch community. I am getting an error Expected object of backend CPU but got backend CUDA for argument #2 ‘weight’. As I searched on internet i have added to(device) to my model and my variables var_X_batch aand var_y_batch . But still i am getting error. Any help ?

import numpy as np
    X=[]
    Y=[]
    for i in range(700):
        X.append(np.load('X_{}.npy'.format(i)))
        Y.append(np.load('Y_{}.npy'.format(i)))
    
    X=np.array(X)
    Y=np.array(Y)

#X.shape is (700,448,448,3)
#y.shape is (700,49,5)

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data
from torch.autograd import Variable

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

def custom_loss(y_pred,y_true):


    for t,p in zip(y_true,y_pred):
        t=t.view(49,5)
        p=p.view(49,5)
        idx=torch.unique(t.nonzero()[:,0])

        xt=t[idx,0]
        xp=p[idx,0]
        yt=t[idx,1]
        yp=p[idx,1]
        eq1=torch.sum((xt-xp)**2 +(yt-yp)**2,dim=0  )
        wt=t[idx,2]
        wp=p[idx,2]
        ht=t[idx,3]
        hp=p[idx,3]
        eq2=torch.sum((torch.sqrt(torch.abs(wt))-torch.sqrt(torch.abs(wp))**2) + (torch.sqrt(torch.abs(ht))-torch.sqrt(torch.abs(hp))**2))

        ct=t[idx,4]
        cp=p[idx,4]
        eq3=torch.sum((ct-cp)**2)
        res=eq1+eq2+10*eq3
        print('res',res)
    return res





BATCH_SIZE = 5

torch_X = torch.from_numpy(X)
torch_X = torch_X.permute(0, 3, 1, 2)

y=np.reshape(Y,(Y.shape[0],-1))
torch_y = torch.from_numpy(y)

# Pytorch train
train = torch.utils.data.TensorDataset(torch_X,torch_y)

# data loader
train_loader = torch.utils.data.DataLoader(train, batch_size = BATCH_SIZE, shuffle = False)



def fit(model, train_loader):
    optimizer = torch.optim.Adam(model.parameters())
    EPOCHS = 50
    model.train()
    
    for epoch in range(EPOCHS):
        correct = 0
        for batch_idx, (X_batch, y_batch) in enumerate(train_loader):
            var_X_batch = X_batch.to(device)
            var_y_batch = y_batch.to(device)
            optimizer.zero_grad()
            output = model(var_X_batch)
            loss = custom_loss(output, var_y_batch)
            loss.backward()
            optimizer.step()


            correct += (output.data == var_y_batch).sum()
            if batch_idx % 50 == 0:
                print('Epoch : {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\t Accuracy:{:.3f}%'.format(
                    epoch, batch_idx*len(X_batch), len(train_loader.dataset), 100.*batch_idx / len(train_loader), loss.data.item(), float(correct*100) / float(BATCH_SIZE*(batch_idx+1))))
                


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        
        self.conv_block = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2) 
        )
        
        self.linear_block = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(128*112*112, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(128, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(64, 245)
        )
        
    def forward(self, x):
        x = self.conv_block(x)
        print(x.shape)
        x = x.view(x.size(0), -1)
        x = self.linear_block(x)
        
        return x

cnn = CNN().to(device)
it = iter(train_loader)
X_batch, y_batch = next(it)
print(cnn.forward(X_batch).shape)


fit(cnn,train_loader)

This line of code will throw the error:

print(cnn.forward(X_batch).shape)

, since you didn’t transferred X_batch to the GPU yet (as is done in your fit method).