Torch("supported devices include CPU, CUDA and HPU, however got MPS")

[LeeeSe]

torch-version: 1.13.1
OS: macOS 13.1
Chip: M1

use tch::{kind, Tensor};

fn main() {
    let model =
        tch::CModule::load_on_device("model/yolov5s.torchscript", tch::Device::Mps).unwrap();
}

[LaurentMazare]

I think the issue is actually on the libtorch side, see pytorch/pytorch#88820
A workaround seems to be to load the model to the cpu and then move it to mps.

[LeeeSe]

Thanks, the problem was solved by the following

fn load_model_to_gpu() {
    let mut vs = tch::nn::VarStore::new(tch::Device::Cpu);
    vs.load("model/yolov5s.torchscript").unwrap();
    println!("vs.device {:?}", vs.device());
    vs.set_device(tch::Device::Mps);
    println!("vs.device now {:?}", vs.device());
}

[guillaume-be]

@LaurentMazare the workaround seems to work for users with access to the VarStore, but causes some frictions when the application does not expose it anymore (e.g. higher level model APIs), see guillaume-be/rust-bert#311

Since this seems to be the only way to load weights on MPS currently, would you be open to a pull request making this the default behaviour when loading weights (i,e. if the VarStore device is Device::Mps, change it to Device::Cpu, load the weights and then change back to Device::Mps)?

This hopefully gets resolved eventually and the workaround could then be removed from the load method.

[LaurentMazare]

@guillaume-be sure, that sounds like a reasonable temporary workaround so feel free to mint a PR or let me know if you don't have the time and I'll do it, probably good to have a comment that points at this issue and mention that we should revisit.