Running with the MPS backend?

[researcx]

https://pytorch.org/docs/stable/notes/mps.html

[ProGamerGov]

@researcx It should be possible to do with a few minor changes I think