Add Apple Silicon MPS support

[HamedMP]

Adds MPS (Meta Performance Shaders) option to run the training on. Torch MPS

The hardcoded "cuda" strings have been changed to use the args.device instead, in case someone wants to try running it on Apple Silicon devices.

P.s. Some formatting has also been done with black formatter.