Fix test_model_parallelism (#25359)

* fix * fix --------- Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
2023-08-08 10:48:45 +02:00
parent d4bd33cc9f
commit 6ea3ee3cd2
25 changed files with 37 additions and 15 deletions
--- a/tests/models/vilt/test_modeling_vilt.py
+++ b/tests/models/vilt/test_modeling_vilt.py
@@ -235,6 +235,7 @@ class ViltModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
    test_pruning = False
    test_headmasking = False
    test_torchscript = False
+    model_split_percents = [0.5, 0.8, 0.9]

    # ViltForMaskedLM, ViltForQuestionAnswering and ViltForImagesAndTextClassification require special treatment
    def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):