Omroth:

 local_labels_float = local_labels.float()


I am not entirely sure what this line is doing in this code here. local_labels_float is not being used at all.



 Omroth:

target_index = Variable(torch.LongTensor([indices[0]]))


Also, why not just transfer target_index to the gpu directly instead of sending labels and then doing max and indexing operations. Maybe that’s the problem.
for batch, labels in training_dataset_generator:
            values, indices = local_labels.max(0)
            target_index = indices[0]
            # transfer to gpu
            local_batch, target_index = batch.to(ml_def.device), target_index.to(ml_def.device)

            prediction = model(local_batch)
            loss = criterion(prediction, target_index)

Note that indices is already a Longtensor.  By the looks of it, you may have to (although I am not entirely sure) use an unsqueeze operation on the target_index, i.e. target_index = target_index.unsqueeze(0) before transferring to the GPU.