torch.load can be unsafe due to pickle exploits

[zwishenzug]

torch.load can be unsafe due to pickle exploits, most projects use a safe loader of some description if they support torch checkpoints.

You could support safetensors also to mitigate that. The code for safe tensors is pretty simple:

from safetensors import safe_open

def safetensors_load(ckpt, map_location="cpu"):
    sd = {}
    with safe_open(ckpt, framework="pt", device=map_location) as f:
        for key in f.keys():
            sd[key] = f.get_tensor(key)
    return {'state_dict': sd}

[ogkalu2]

@zwishenzug Thanks for the help. I'll implement them all shortly