Hi,
I have a model saved with torch.save into the file. The file is quite big (say, 100 GB), torch.load is crashing on a nodes with moderate CPU RAM. The good part is that I don’t need all the tensors in CPU memory at once.
Is there a way to customize torch.load so that it doesn’t produce model with all tensors deserialized? Basically, the majority of tensors should be dropped right after each of them being deserialized or not deserialized at all.
The docs suggest 2 map_location callable argument could possibly do the trick, I cannot find examples for that though.
Any help is appreciated!
Thanks,
Max.