Running nnUNet_train on M2 Mac

[vcasellesb]

Hi FabianIsensee, I was trying to run training on nnUNet using a Macbook Pro with Apple Silicon. I was wondering, is there any way of instead of using CUDA using the MPS (Mac GPU backend) for speeding up the training?

Just running nnUNet_train uses CPU instead, but I would like it to also use the built-in GPU in Apple Chips. Would there be any relatively quick fix?

Thank you very much!!

Best,

Vicent Caselles

[FabianIsensee]

Hi Vincent,
I don't know how to do that. Never worked on a Mac in my life.
You probably need to install a pytorch version that is compatible with the Apple hardware and then (maybe?) change the device in the code from cuda to whatever the mac one is called.
Best,
Fabian

[vcasellesb]

Thank you very much Fabian. I've been trying but it seems Pytorch isn't fully stable on Mac devices. Will probably run it on a server with CUDA.

Best,

Vicent

[FabianIsensee]

Sorry I couldn't help you :-/ Let's hope it becomes stable soon!