Fix: Update torch.load to use weights_only=True to prevent security warning

This commit is contained in:
yaslack 2024-11-20 18:05:08 +01:00
parent 173ff7dd1d
commit 83876b4bc7

View File

@ -147,7 +147,7 @@ def load_model(
with (
io.BytesIO(checkpoint_file) if in_memory else open(checkpoint_file, "rb")
) as fp:
checkpoint = torch.load(fp, map_location=device)
checkpoint = torch.load(fp, map_location=device, weights_only=True)
del checkpoint_file
dims = ModelDimensions(**checkpoint["dims"])