Give more memory in test_disk_offload (#25315)

This commit is contained in:
Sylvain Gugger 2023-08-04 14:10:31 +02:00 committed by GitHub
parent 67683095a6
commit fab1a0aa82
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -2521,7 +2521,7 @@ class ModelTesterMixin:
base_output = model(**inputs_dict_class)
model_size = compute_module_sizes(model)[""]
max_size = int(self.model_split_percents[0] * model_size)
max_size = int(self.model_split_percents[1] * model_size)
with tempfile.TemporaryDirectory() as tmp_dir:
model.cpu().save_pretrained(tmp_dir)