Update __init__.py

This commit is contained in:
Jong Wook Kim 2025-06-25 17:51:51 -07:00 committed by GitHub
parent 915b3e747b
commit 0692eb6b29
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -147,10 +147,8 @@ 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:
try: kwargs = {"weights_only": True} if torch.__version__ >= "1.13" else {}
checkpoint = torch.load(fp, map_location=device, weights_only=True) checkpoint = torch.load(fp, map_location=device, **kwargs)
except TypeError: # for compatibility with older torch
checkpoint = torch.load(fp, map_location=device)
del checkpoint_file del checkpoint_file
dims = ModelDimensions(**checkpoint["dims"]) dims = ModelDimensions(**checkpoint["dims"])