xpected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

def accuracy(predictions, labels):
    classes = torch.argmax(predictions, dim=1)
    print("prediction",classes)
    print("ground truth",labels)
    return torch.mean((classes == labels).float())

def test_loop(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            print(pred.argmax(1))
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

if __name__ == '__main__':
    "Image Visualization"
    # iter_l = iter(train_loader)
    # input, classes = next(iter_l)
    # out = torchvision.utils.make_grid(input)
    # print(out.shape)
    # print([x.item() for x in classes])
    # imshow(out, title=[x.item() for x in classes])
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    "Creating CoAtNet Model"
    model = timm.create_model("coatnet_2_rw_224.sw_in12k_ft_in1k", pretrained=True, num_classes=1981, scriptable=True,
                              exportable=True)
    model = torch.nn.DataParallel(model)
    model.to(device)
    model.cuda()
    print(device)
    # print(model.get_classifier())
    data_config = timm.data.resolve_model_data_config(model)
    print(data_config)
    "Creating optimizer"
    optimizer = timm.optim.create_optimizer_v2(model, opt='adam', lr=0.0001)
    "creating Loss function"
    loss_fn = torch.nn.CrossEntropyLoss()
    for epoch in range(1):
        running_loss = 0.00
        running_accuracy = 0.00
        size = len(train_loader.dataset)
        for batch,(inputs, targets) in enumerate(train_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = loss_fn(outputs, targets)
            # print(loss)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            
            if batch % 100 == 0: 
                loss, current = loss.item(), batch * len(inputs)
                # running_loss += loss.item()
                running_accuracy += accuracy(outputs, targets)
                print("Accuracy",running_accuracy)
                print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
        # running_loss /= len(trainloader)
        running_accuracy /= len(train_loader)
        test_loop(val_loader, model,loss_fn)
        
    model = model.eval()
    # saving the entire model
    torch.save(model,"test_coatnet.pt")

You are not pushing the data to the GPU in your test_loop:

for X, y in dataloader:
    pred = model(X)

which is most likely causing the issue.

1 Like

Thank you It worked!