[Misc] Fix issues reported by torchfix (#4837)

This commit is contained in:
Brayden Zhong
2025-03-27 23:10:32 -04:00
committed by GitHub
parent b149b39353
commit e84f4ba0ab
3 changed files with 5 additions and 5 deletions

View File

@@ -92,7 +92,7 @@ def convert_bin_to_safetensor_file(
pt_filename: str,
sf_filename: str,
) -> None:
loaded = torch.load(pt_filename, map_location="cpu")
loaded = torch.load(pt_filename, map_location="cpu", weights_only=True)
if "state_dict" in loaded:
loaded = loaded["state_dict"]
shared = _shared_pointers(loaded)
@@ -380,7 +380,7 @@ def np_cache_weights_iterator(
disable=not enable_tqdm,
bar_format=_BAR_FORMAT,
):
state = torch.load(bin_file, map_location="cpu")
state = torch.load(bin_file, map_location="cpu", weights_only=True)
for name, param in state.items():
param_path = os.path.join(np_folder, name)
with open(param_path, "wb") as f: