Add support for device in safetensors.torch.load_model

[Wauplin]

safetensors.torch.load_file has a "device" parameter to load the tensors directly to the correct device. This PR adds support for this parameter in safetensors.torch.load_model too.

(also fix device type -see #449 (comment))

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[mishig25]

Suggested change

	device (`Dict[str, any]`, *optional*, defaults to `cpu`):
	device (`str`, *optional*, defaults to `cpu`):

is it dict or str ?

[Wauplin]

I took the exact same definition as in load_file docstring:

safetensors/bindings/python/py_src/safetensors/torch.py

Line 291 in b947b59

device (`Dict[str, any]`, *optional*, defaults to `cpu`):

[mishig25]

defaults to cpu

I see. I guess it is either str or Union[Dict[str, any], str]

[Wauplin]

Made the change in d723ff7 :) (and used the Union[Dict[str, any], str] type).

[Wauplin]

Go down to the rust code here:

#[derive(Debug, Clone, PartialEq, Eq)]
enum Device {
    Cpu,
    Cuda(usize),
    Mps,
    Npu(usize),
    Xpu(usize),
}

[mishig25]

why is it Union[Dict[str, any], str] rather than str ?

Doesn't the enum mean that value is str rather than dict?

enum Device {
    Cpu,
    Cuda(usize),
    Mps,
    Npu(usize),
    Xpu(usize),
}

[Wauplin]

Yes sorry so I went back to the source code and it looks like the expected type should be Union[str, int]. Here is the code:

• expecting string
• expecting int (which has been discussed/added in Some PyO3 specifics. #35 (comment)) => defaults to "Cuda:<int_value>".

I have updated the types accordingly in 3bad1e2. Let's wait for @Narsil's return just to be sure.

[Narsil]

LGTM