From d291b052db19e4c2a27f76eb556d2c87f5b572fa Mon Sep 17 00:00:00 2001 From: vigneshkumar25 Date: Wed, 1 Jul 2026 19:12:47 +0530 Subject: [PATCH] fix: support transformers>=5.0.0 and fix torch.load security warning - Relax transformers upper bound to <6.0 to unblock v5 users - Add weights_only=True to torch.load to fix CVE security warning Fixes #2432 --- optimum/fx/parallelization/utils.py | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/optimum/fx/parallelization/utils.py b/optimum/fx/parallelization/utils.py index e1d6120628..9c6a4f0233 100644 --- a/optimum/fx/parallelization/utils.py +++ b/optimum/fx/parallelization/utils.py @@ -445,7 +445,7 @@ def convert_bin_to_safetensors( output_dir = cache_dir if cache_dir else weight_file_path.parent output_file_path = os.path.join(output_dir, safetensors_filename) if not os.path.isfile(output_file_path): - checkpoint = torch.load(weight_file, map_location=torch.device("cpu")) + checkpoint = torch.load(weight_file, map_location=torch.device("cpu"), weights_only=True) data_pointers = set() for k, v in checkpoint.items(): if v.data_ptr() in data_pointers: diff --git a/setup.py b/setup.py index bf9edea76d..ad11ab1938 100644 --- a/setup.py +++ b/setup.py @@ -13,7 +13,7 @@ REQUIRED_PKGS = [ - "transformers>=4.29", + "transformers>=4.29,<6.0", "torch>=1.11", "packaging", "numpy",