Allocate model and data on GPU

I shift my data and model on GPU but still it gives this error.

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument target in method wrapper_nll_loss_forward)

model = ConvNet().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

def batch_gd(model, criterion, optimizer, train_loader, test_loader, epochs):
train_losses = np.zeros(epochs)
test_losses = np.zeros(epochs)

for it in range(epochs):
    t0 = datetime.datetime.now()
    train_loss = []
    for sets, labels in train_loader:
        labels = labels.type(torch.LongTensor)
        sets, targets = sets.to(device), labels.to(device)
    
        # Zero the parameter gradientd
        optimizer.zero_grad()

        # Forward pass
        outputs = model(sets)
        loss = criterion(outputs, labels)

        # Backwords and optimize
        loss.backward()
        optimizer.step()

        train_loss.append(loss.item())

    # Get train loss and test loss
    train_loss = np.mean(train_loss) 

    test_loss = []
    for sets, labels in test_loader:
        labels = labels.type(torch.LongTensor)
        sets, labels = sets.to(device), labels.to(device)
        outputs = model(sets)
        loss = criterion(outputs, labels)
        test_loss.append(loss.item())
        
    test_loss = np.mean(train_loss)
    
    # Save losses
    train_losses[it] = train_loss
    test_losses[it] = test_loss
    
    dt = datetime.datetime.now() - t0
    print(f'Epoch {it+1}/{epochs}, Train Loss: {train_loss:.4f}, \n Test Loss: {test_loss:.4f}, Duration: {dt}')
    
return train_losses, test_losses

You are pushing the labels tensor to the GPU and assign it to targets here:

sets, targets = sets.to(device), labels.to(device)

However, later you are using labels (the CPU tensor) which yields this error:

loss = criterion(outputs, labels)

Use targets or reassign labels to itself.

1 Like

Thanks for your reply, this is my silly mistake. I am new in PyTorch.