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()