Step towards adding MPS support.

[anuragg1209]

Hi @LeoGrin & @noahho,

This PR is a step towards adding support for Apple Silicon (mps) #269 . I believe the mps support will not work directly. The two main issues with mps are that:

1. Torch does not support autocast for mps yet.
2. The MPS framework doesn't support float64 tensors. We have to use float32 for it to function properly. Due to this, we need to make changes in other areas. I would like to hear your thoughts on this. Please correct me if I am mistaken.

[LeoGrin]

Thanks @anuragg1209!

The MPS framework doesn't support float64 tensors. We have to use float32 for it to function properly. Due to this, we need to make changes in other areas. I would like to hear your thoughts on this. Please correct me if I am mistaken.

Aren't we defaulting to float32 for the pytorch part (using 64 for the sklearn part only?)? Not sure about that, but I think we can (should?) do that if we need to.

I see that the CI tests on macbook are failing with a memory error. It makes sense that we're using more memory if we're not using autocast, but it seems worrying as the test examples are very lightweight. Do you have an idea of what is the issue? Perhaps one of:

1. There is an issue with the memory estimation?
2. Our model/code takes too much memory for a standard mps setting
3. The CI hardware is unrealistically weak