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 ( with (
io.BytesIO(checkpoint_file) if in_memory else open(checkpoint_file, "rb") io.BytesIO(checkpoint_file) if in_memory else open(checkpoint_file, "rb")
) as fp: ) as fp:
checkpoint = torch.load(fp, map_location=device) checkpoint = torch.load(fp, map_location=device, weights_only=True)
del checkpoint_file del checkpoint_file
dims = ModelDimensions(**checkpoint["dims"]) dims = ModelDimensions(**checkpoint["dims"])