Closed sameastham closed 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!
Closing this issue for now. Reopen if you have further issues.
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