From a7cb92aa030bc4dde3ea82aad849e783d15f470a Mon Sep 17 00:00:00 2001
From: Yih-Dar <2521628+ydshieh@users.noreply.github.com>
Date: Fri, 2 Feb 2024 14:11:50 +0100
Subject: [PATCH] fix / skip (for now) some tests before switch to torch 2.2
 (#28838)

* fix / skip some tests before we can switch to torch 2.2

* style

---------

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
---
 tests/models/mega/test_modeling_mega.py | 13 +++++++++++++
 tests/models/vits/test_modeling_vits.py |  5 +++++
 tests/utils/test_model_output.py        |  4 +++-
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/tests/models/mega/test_modeling_mega.py b/tests/models/mega/test_modeling_mega.py
index a67ee0d003..e07f4efedd 100644
--- a/tests/models/mega/test_modeling_mega.py
+++ b/tests/models/mega/test_modeling_mega.py
@@ -19,6 +19,7 @@ import unittest
 from transformers import MegaConfig, is_torch_available
 from transformers.testing_utils import (
     TestCasePlus,
+    is_flaky,
     require_torch,
     require_torch_fp16,
     slow,
@@ -534,6 +535,18 @@ class MegaModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin
         self.model_tester = MegaModelTester(self)
         self.config_tester = ConfigTester(self, config_class=MegaConfig, hidden_size=37)
 
+    # TODO: @ydshieh
+    @is_flaky(description="Sometimes gives `AssertionError` on expected outputs")
+    def test_pipeline_fill_mask(self):
+        super().test_pipeline_fill_mask()
+
+    # TODO: @ydshieh
+    @is_flaky(
+        description="Sometimes gives `RuntimeError: probability tensor contains either `inf`, `nan` or element < 0`"
+    )
+    def test_pipeline_text_generation(self):
+        super().test_pipeline_text_generation()
+
     def test_config(self):
         self.config_tester.run_common_tests()
 
diff --git a/tests/models/vits/test_modeling_vits.py b/tests/models/vits/test_modeling_vits.py
index abae6999cd..c1c4117f7e 100644
--- a/tests/models/vits/test_modeling_vits.py
+++ b/tests/models/vits/test_modeling_vits.py
@@ -176,6 +176,11 @@ class VitsModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
     def test_config(self):
         self.config_tester.run_common_tests()
 
+    # TODO: @ydshieh
+    @is_flaky(description="torch 2.2.0 gives `Timeout >120.0s`")
+    def test_pipeline_feature_extraction(self):
+        super().test_pipeline_feature_extraction()
+
     @unittest.skip("Need to fix this after #26538")
     def test_model_forward(self):
         set_seed(12345)
diff --git a/tests/utils/test_model_output.py b/tests/utils/test_model_output.py
index 33013f222e..b1fc3dd013 100644
--- a/tests/utils/test_model_output.py
+++ b/tests/utils/test_model_output.py
@@ -155,9 +155,11 @@ class ModelOutputTester(unittest.TestCase):
         if is_torch_greater_or_equal_than_2_2:
             self.assertEqual(
                 pytree.treespec_dumps(actual_tree_spec),
-                '[1, {"type": "tests.utils.test_model_output.ModelOutputTest", "context": ["a", "c"], "children_spec": [{"type": null, "context": null, "children_spec": []}, {"type": null, "context": null, "children_spec": []}]}]',
+                '[1, {"type": "tests.utils.test_model_output.ModelOutputTest", "context": "[\\"a\\", \\"c\\"]", "children_spec": [{"type": null, "context": null, "children_spec": []}, {"type": null, "context": null, "children_spec": []}]}]',
             )
 
+    # TODO: @ydshieh
+    @unittest.skip("CPU OOM")
     @require_torch
     def test_export_serialization(self):
         if not is_torch_greater_or_equal_than_2_2: