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