Skip to content

Commit c60aa95

Browse files
applesaucethebunb8zhong
authored andcommitted
fix: apply torchfix weights_only suggestion
1 parent a5337f5 commit c60aa95

File tree

2 files changed

+2
-2
lines changed

2 files changed

+2
-2
lines changed

tools/checkpoint_conversion/convert_roberta_checkpoints.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ def convert_checkpoints(size):
6060
checkpoint_path = os.path.join(extract_dir, "model.pt")
6161

6262
# Load PyTorch RoBERTa checkpoint.
63-
pt_ckpt = torch.load(checkpoint_path, map_location=torch.device("cpu"))
63+
pt_ckpt = torch.load(checkpoint_path, map_location=torch.device("cpu"), weights_only=True)
6464
pt_cfg = pt_ckpt["args"]
6565
pt_model = pt_ckpt["model"]
6666

tools/checkpoint_conversion/convert_xlm_roberta_checkpoints.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ def convert_checkpoints(size):
4545
checkpoint_path = os.path.join(extract_dir, "model.pt")
4646

4747
# Load PyTorch XLM-R checkpoint.
48-
pt_ckpt = torch.load(checkpoint_path, map_location=torch.device("cpu"))
48+
pt_ckpt = torch.load(checkpoint_path, map_location=torch.device("cpu"), weights_only=True)
4949
pt_cfg = pt_ckpt["args"]
5050
pt_model = pt_ckpt["model"]
5151

0 commit comments

Comments
 (0)