Set weights_only in torch.load (#36991)

2025-03-27 22:55:50 +08:00
parent de77f5b1ec
commit 41a0e58e5b
28 changed files with 64 additions and 78 deletions
--- a/examples/legacy/token-classification/utils_ner.py
+++ b/examples/legacy/token-classification/utils_ner.py
@@ -242,7 +242,7 @@ if is_torch_available():
            with FileLock(lock_path):
                if os.path.exists(cached_features_file) and not overwrite_cache:
                    logger.info(f"Loading features from cached file {cached_features_file}")
-                    self.features = torch.load(cached_features_file)
+                    self.features = torch.load(cached_features_file, weights_only=True)
                else:
                    logger.info(f"Creating features from dataset file at {data_dir}")
                    examples = token_classification_task.read_examples_from_file(data_dir, mode)