Update loss.py

[Luccifer]

Fixed an issue of RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! According to the https://github.com/andrewjong/SwapNet/issues/47 - the issue still exists. Managed to run it with a little fix