Make it available on GPU device

[JosieHong]

Hi,
I made the following modifications to make this loss function available on both CPU and GPU devices:

1. Check the device of labels and logits;
2. Define the device;
3. Move weights to that device.

Thanks,
Josie

[fcakyon]

@JosieHong feel free to use this improved/fixed/battle-tested version, includes auto device selection: https://github.com/fcakyon/balanced-loss#improvements

[JosieHong]

@JosieHong feel free to use this improved/fixed/battle-tested version, includes auto device selection: https://github.com/fcakyon/balanced-loss#improvements

Thanks, @fcakyon! I had been struggling with a proper implementation for a long time. This implementation looks very nice, and I am excited to try it!