Merge bf24f74a349445487d53f58cbacc8d17fdb10c38 into c0d2f624c09dc18e709e37c2ad90c039a4eb72a2

This commit is contained in:
R I P Skillan 2025-09-27 01:32:11 +00:00 committed by GitHub
commit 7a48d668b5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 3 additions and 1 deletions

View File

@ -5,3 +5,4 @@ tqdm
more-itertools
tiktoken
triton>=2.0.0;platform_machine=="x86_64" and sys_platform=="linux" or sys_platform=="linux2"
packaging

View File

@ -4,6 +4,7 @@ import os
import urllib
import warnings
from typing import List, Optional, Union
from packaging import version
import torch
from tqdm import tqdm
@ -147,7 +148,7 @@ def load_model(
with (
io.BytesIO(checkpoint_file) if in_memory else open(checkpoint_file, "rb")
) as fp:
kwargs = {"weights_only": True} if torch.__version__ >= "1.13" else {}
kwargs = {"weights_only": True} if version.parse(torch.__version__ ) >= version.parse("1.13") else {}
checkpoint = torch.load(fp, map_location=device, **kwargs)
del checkpoint_file