FEATURE: mps support

[michaelbarkasi]

Feature you'd like to see:

The pytorch code is written to call torch.cuda explicitly. It will not try to call mps if 'device = torch.device("mps")' is passed into the functions. It only took me a few minutes and a few modified lines to get kilosort4 running on mps. It seems like this would be a pretty easy thing to implement.

Additional Context

No response

[jacobpennington]

Would you mind sharing which lines you modified? The code is already intended to use whatever device you specify.

[michaelbarkasi]

Thanks for the fast reply! These are the primary lines, although I spoke a little too soon. There are other calls throughout the code that are incompatible with mps, e.g., right now I'm working on addressing the calls to coo.

[jacobpennington]

Thanks. Would you be able to submit a pull request with the changes when you're finished?

[michaelbarkasi]

yes!

[michaelbarkasi]

After digging more, there's actually quite a lot that needs revision. The main issue is that Pytorch's support for mps is more limited than for cuda, and there are several torch functions/methods called (e.g., coo, but others) which are not supported when torch tensors are on mps. For now I'm putting this on the shelf, but if I come back to it and finish I'll put in a pull request.

[michaelbarkasi]

eh... okay, I spoke too soon (again). It didn't take that much more work. I believe everything is working now on mps. I submitted a pull request.

[jacobpennington]

Thanks again for working on this, I'm going to close the issue for now and move the conversation to the related pull request. For anyone looking into this feature in the future, see #839