MPS Load capabilities

[guillaume-be]

• Temporary fix to the VarStore load method allowing loading weights on MPS-accelerated machines
  The device of the VarStore is temporarily set to CPU when loading the weights and then changed back to MPS as described by the workarounds described in Add MPS as default if available guillaume-be/rust-bert#311 and Torch("supported devices include CPU, CUDA and HPU, however got MPS") #609

I do not have a MPS device to test the behaviour. It would be great if one of the contributors who identified the issue and the workaround could validate it solves the loading issues before merging.

[bakkot]

I am not one of the original contributors, but I do have an MPS device, and I confirm this fixed the issue.

Specifically, I built https://github.com/LaurentMazare/diffusers-rs with the device manually set to MPS by changing this line to Device::Mps. When specifying this crate as a dependency at main, that fails with the error mentioned in #609. When I use this crate at this PR, it works, and runs several times faster than the non-MPS version.

While I'm here, I note that this PR doesn't restore the original device in the two error cases (the ? and the explicit return Err). That seems like it might be confusing.

[LaurentMazare]

Thanks for confirming, agreed that it would be nicer to set the device back to Mps if an error causes an early exit, @guillaume-be is it something that you could look at and then I think this is good to be merged.

[guillaume-be]

Thank you very much for confirming @bakkot - @LaurentMazare yes I should be able to implement a more robust setup later today

[guillaume-be]

@bakkot , @LaurentMazare I have wrapped the tensor loading in a closure to ensure the device reset code gets executed before the closure Result gets passed on to the caller

[LaurentMazare]

Looks good, thanks for the PR!