Some PyTorch operations are not implemented in MPS yet and will throw an error. To avoid this, you should set the environment variable PYTORCH_ENABLE_MPS_FALLBACK=1 to use the CPU kernels instead (you'll still see a UserWarning).

If you run into any other errors, please open an issue in the PyTorch repository because the [Trainer] only integrates the MPS backend.

With the mps device set, you can: