huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 12:24:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

Force weights_only (before fully breaking pickle files anyway).

Browse Source
This commit is contained in:
Nicolas Patry 2024-04-05 16:17:16 +00:00
parent f9958ee191
commit 422f23be74
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
server/text_generation_server/utils/convert.py
View File

@ -68,7 +68,7 @@ def convert_file(pt_file: Path, sf_file: Path, discard_names: List[str]):
Forcing us to check for potentially different keys during load when looking Forcing us to check for potentially different keys during load when looking
for specific tensors (making tensor sharing explicit). for specific tensors (making tensor sharing explicit).
""" """
loaded = torch.load(pt_file, map_location="cpu") loaded = torch.load(pt_file, map_location="cpu", weights_only=True)
if "state_dict" in loaded: if "state_dict" in loaded:
loaded = loaded["state_dict"] loaded = loaded["state_dict"]
to_removes = _remove_duplicate_names(loaded, discard_names=discard_names) to_removes = _remove_duplicate_names(loaded, discard_names=discard_names)