search

MIC-DKFZ / nnUNet

Apache License 2.0
5.9k stars 1.76k forks source link

mps not supported #2244

Closed tpecot closed 5 months ago

tpecot commented 5 months ago

Hi! I installed nnUNet and try to train a model. I have a macBook pro with a M3 max chip. I can train the model with the CPU, I tried to train it with the mps device but I get this:

torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
AssertionError: Device mps not supported

Apparently, the inductor backend is not supported on MPS: https://discuss.pytorch.org/t/jitting-fails-on-mps-device/192079 still the case in April: https://discuss.pytorch.org/t/questions-about-pytorch-2-0-and-mps/168552 Do you have any idea how to deal with that? I guess I could download your code and change the backend to aot_eager but I'd rather not if possible. Thanks!

Karol-G commented 5 months ago

Hey @tpecot,

My apologies but we do not provide support for systems other than Linux. However, feel free to share here any insights you gain in your setup process for Mac OS which might help other users with Mac OS.

Best regards, Karol

tpecot commented 5 months ago

Hi @Karol-G, thanks for your help! I'll check by myself then. Maybe a suffestion, the documentation says that it's possible to use mps device by adding: -device mps Maybe it'd be useful to users to add that it does not work.

Best, Thierry

ancestor-mithril commented 5 months ago

You can set the environment variable nnUNet_compile=f before running nnUNet and it should also work on mps.

tpecot commented 5 months ago

Thanks @ancestor-mithril! I'll check it when I'm back home.

tpecot commented 5 months ago

It actually works, thanks!!! Maybe you could add a comment in the documentation when talking about mps device, add a sentence to state that you need to change the envirnoment variable nnUNet_compile=f, that could be helpful for other.

FabianIsensee commented 5 months ago

We now just disable torch.compile for mps devices. No user action required