Fix more offload edge cases (#25342)

* fix * fix * fix --------- Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
2023-08-07 17:45:41 +02:00
parent 7d65697da7
commit c177606fb4
2 changed files with 8 additions and 3 deletions
--- a/tests/models/longformer/test_modeling_longformer.py
+++ b/tests/models/longformer/test_modeling_longformer.py
@@ -326,6 +326,9 @@ class LongformerModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCa
        else {}
    )

+    # Need to use `0.6` instead of `0.5` for `test_disk_offload`
+    model_split_percents = [0.6, 0.7, 0.9]
+
    # TODO: Fix the failed tests
    def is_pipeline_test_to_skip(
        self, pipeline_test_casse_name, config_class, model_architecture, tokenizer_name, processor_name