SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Set weights_only in torch.load (#36991)

Browse Source
This commit is contained in:
cyyever
2025-03-27 22:55:50 +08:00
committed by GitHub
parent de77f5b1ec
commit 41a0e58e5b
28 changed files with 64 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/models/patchtst/test_modeling_patchtst.py
View File

@@ -303,7 +303,7 @@ class PatchTSTModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase
def prepare_batch(repo_id="hf-internal-testing/etth1-hourly-batch", file="train-batch.pt"):
file = hf_hub_download(repo_id=repo_id, filename=file, repo_type="dataset")
batch = torch.load(file, map_location=torch_device)
batch = torch.load(file, map_location=torch_device, weights_only=True)
return batch