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/autoformer/test_modeling_autoformer.py
View File

@@ -415,7 +415,7 @@ class AutoformerModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCa
def prepare_batch(filename="train-batch.pt"):
file = hf_hub_download(repo_id="hf-internal-testing/tourism-monthly-batch", filename=filename, repo_type="dataset")
batch = torch.load(file, map_location=torch_device)
batch = torch.load(file, map_location=torch_device, weights_only=True)
return batch

2
tests/models/idefics/test_image_processing_idefics.py
View File

@@ -28,7 +28,7 @@ if is_torch_available():
import torch
if is_torchvision_available():
import torchvision.transforms as transforms
from torchvision import transforms
if is_vision_available():
from PIL import Image

2
tests/models/informer/test_modeling_informer.py
View File

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

4
tests/models/llava_next/test_modeling_llava_next.py
View File

@@ -408,7 +408,7 @@ class LlavaNextForConditionalGenerationIntegrationTest(unittest.TestCase):
filename="llava_1_6_input_ids.pt",
repo_type="dataset",
)
original_input_ids = torch.load(filepath, map_location="cpu")
original_input_ids = torch.load(filepath, map_location="cpu", weights_only=True)
# replace -200 by image_token_index (since we use token ID = 32000 for the image token)
# remove image token indices because HF impl expands image tokens `image_seq_length` times
original_input_ids = original_input_ids[original_input_ids != -200]
@@ -420,7 +420,7 @@ class LlavaNextForConditionalGenerationIntegrationTest(unittest.TestCase):
filename="llava_1_6_pixel_values.pt",
repo_type="dataset",
)
original_pixel_values = torch.load(filepath, map_location="cpu")
original_pixel_values = torch.load(filepath, map_location="cpu", weights_only=True)
assert torch.allclose(original_pixel_values, inputs.pixel_values.half())
# verify generation

2
tests/models/patchtsmixer/test_modeling_patchtsmixer.py
View File

@@ -452,7 +452,7 @@ class PatchTSMixerModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.Test
def prepare_batch(repo_id="ibm/patchtsmixer-etth1-test-data", file="pretrain_batch.pt"):
# TODO: Make repo public
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

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

2
tests/models/time_series_transformer/test_modeling_time_series_transformer.py
View File

@@ -481,7 +481,7 @@ class TimeSeriesTransformerModelTest(ModelTesterMixin, PipelineTesterMixin, unit
def prepare_batch(filename="train-batch.pt"):
file = hf_hub_download(repo_id="hf-internal-testing/tourism-monthly-batch", filename=filename, repo_type="dataset")
batch = torch.load(file, map_location=torch_device)
batch = torch.load(file, map_location=torch_device, weights_only=True)
return batch

2
tests/models/videomae/test_modeling_videomae.py
View File

@@ -456,7 +456,7 @@ class VideoMAEModelIntegrationTest(unittest.TestCase):
# add boolean mask, indicating which patches to mask
local_path = hf_hub_download(repo_id="hf-internal-testing/bool-masked-pos", filename="bool_masked_pos.pt")
inputs["bool_masked_pos"] = torch.load(local_path)
inputs["bool_masked_pos"] = torch.load(local_path, weights_only=True)
# forward pass
with torch.no_grad():