Make sure that logits are always on the device where the model weights are

[ManilShrestha]

The issue was due to the logits always being in device=0, I made it such that the logits generated by the model should always be in the model's device.

[rlouf]

What happens when the model is loaded on several GPUs? I think that was the rationale behind pinning the logits to device 0.

[rlouf]

Note that this is going to become irrelevant once #926 is merged.