In some cases there are operations not supported on mps. For those set:

os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

and it will run on cpu if some operation isn't supported