I’m trying to train with CrossEntropyLoss but my argmax method is failing with
RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 ‘target’ in call to _thnn_nll_loss_forward
at the loss = criterion(prediction, target_index) line. Can anyone point to my mistake?
training_dataset_generator = torch.utils.data.DataLoader(traning_dataset, **params)
criterion = torch.nn.CrossEntropyLoss(reduction='sum', weight=torch.Tensor([1.0, 10.0]))
optimizer = torch.optim.SGD(model.parameters(), lr=ml_def.learning_rate)
for epoch in range(max_epochs):
print("doing epoch")
epoch_loss = 0.0
for batch, labels in training_dataset_generator:
# transfer to gpu
local_batch, local_labels = batch.to(ml_def.device), labels.to(ml_def.device)
prediction = model(local_batch)
local_labels_float = local_labels.float()
values, indices = local_labels.max(0)
target_index = Variable(torch.LongTensor([indices[0]]))
loss = criterion(prediction, target_index)
epoch_loss += loss
#print("val", local_batch[0], "original", labels[0], "prediction", prediction.detach()[0], "loss", loss)
optimizer.zero_grad()
loss.backward()
optimizer.step()