consider cpu fallback for apple silicon devices #477
Description
pytorch does not support all operations on the accelerated backend on macos "mps" (metal performance shader). The safe alternative is running things on the cpu backend - but that's slow. pytorch provides an environment variable that should enable always using the "mps" backend and automatically falling back to cpu for unsupported operations. In ilastik we have the following snippet, which starts a suprocess running inference with a modified env:
@property
def _proc_env(self):
"""Modified env for tiktorch server subprocess"""
env = os.environ.copy()
if platform.system().lower() == "darwin" and platform.machine().lower() == "arm64":
# pytorch versions up to the 2.6 don't support all operations (esp 3d) on MPS
# this env variable allows falling back to CPU for those networks instead of failing
# see for current status https://github.com/pytorch/pytorch/issues/141287
if "PYTORCH_ENABLE_MPS_FALLBACK" not in env:
env["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
logging.info(f"Added 'PYTORCH_ENABLE_MPS_FALLBACK=1' to tiktorch env.")
return env