Compatibility with torch.save()?

[dconathan]

Very cool project!

I was wondering if there's possibility to use this package as the "backend" for torch.save to enable easier integration with downstream projects that are already using torch.save/load...

seems like it might be possible since torch.save and torch.load have a pickle_module= argument. Something like:

import safetensors
torch.save(my_tensor, "my_tensor.pt", pickle_module=safetensors)   # or safetensors.torch

[Narsil]

Hi @dconathan ,

No it's not possible nor intended.
Pytorch uses pickle + zip.

If you care only about security, there's a new weights_only=True parameter https://pytorch.org/docs/stable/generated/torch.load.html which you can use without changing too much existing codebases.

If you want stuff like lazy-loading (only part of tensors on files) or the speedups I mention, then you cannot (afaik) use torch.load, torch.save.

But you can use safetensors layouts without even using this library:
https://gist.github.com/Narsil/3edeec2669a5e94e4707aa0f901d2282

[dconathan]

Thanks for the reply! Yeah I looked a little into the src of torch.save and doesn't seem possible... thanks for the links.

[julien-c]

this gist your linked is very cool @Narsil!

[Narsil]

It's exactly what this lib does for CPU, minus handling various torch versions and falling back to slower alternatives.