hukkelas / deep_privacy2

DeepPrivacy2 - A Toolbox for Realistic Image Anonymization
Apache License 2.0
292 stars 33 forks source link

MPS on MacOS? #12

Closed sameastham closed 1 year ago

sameastham commented 1 year ago

Before I start digging through the code, is there a straightforward way to use MPS for M1/M2 Macs rather than CUDA for NVIDIA? PyTorch does support MPS: https://pytorch.org/docs/stable/notes/mps.html

hukkelas commented 1 year ago

I have not tested with MPS on MacOS and it is not supported out of the box. The framework assumes CUDA device.

All CPU-to-CUDA transfers are done with helper function from the extension "tops". To use MPS instead you should only need to change the following functions: to_cuda and get_device.

Let me know if you get it to work!

hukkelas commented 1 year ago

Closing this issue for now. Reopen if you have further issues.