safetensors

safetensors copied to clipboard

The code in get_tensor assumes that torch.asarray will always load the tensor in cpu, and then the tensor is copied to self.device in the "to" call at the end of the function. However, that assumption is not always correct. To make sure, device=cpu has to be passed to asarray.

Another option would be to set the final device since the beginning:

let device = match self.device {
    Device::Cpu     => format!("cpu"),
    Device::Cuda(n) => format!("cuda:{n}"),
    Device::Mps     => format!("mps"),
    Device::Npu(n)  => format!("npu:{n}"),
    Device::Xpu(n)  => format!("xpu:{n}"),
}.into_py(py);

let kwargs = [(intern!(py, "dtype"), torch_uint8),(intern!(py, "device"), device)].into_py_dict(py);

I'm not sure what the rationale is for loading first into CPU and then moving to self.device at the end.

Fixes issue #437

@Narsil, I'm wondering of any of the maintainers could take a look at this. It's a simple issue, but it's blocking some of our users. Thanks!

Any updates on this?

No, I haven't heard back from the maintainers on this.