From 422f23be74da441bca3169ab2ae66460ba580978 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Fri, 5 Apr 2024 16:17:16 +0000 Subject: [PATCH] Force weights_only (before fully breaking pickle files anyway). --- server/text_generation_server/utils/convert.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/text_generation_server/utils/convert.py b/server/text_generation_server/utils/convert.py index 0b62f520..d9c3276b 100644 --- a/server/text_generation_server/utils/convert.py +++ b/server/text_generation_server/utils/convert.py @@ -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 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: loaded = loaded["state_dict"] to_removes = _remove_duplicate_names(loaded, discard_names=discard_names)