Device error when extending Trainer class

[sgugger]

@gianluca-db0

loss_fct = nn.CrossEntropyLoss(weight=torch.tensor([0.2, 0.3, 0.5]))

in this line the weights should be put on the same device as the model.