From acc439ba172027152fb394c28b6f55a8df785326 Mon Sep 17 00:00:00 2001
From: Arthur <48595927+ArthurZucker@users.noreply.github.com>
Date: Tue, 6 Dec 2022 16:14:03 +0100
Subject: [PATCH] Ci-jukebox (#20613)

* fix cuda OOM by using single Prior

* only send to device when used

* use custom model

* Skip the big slow test

* Update tests/models/jukebox/test_modeling_jukebox.py

Co-authored-by: Yih-Dar <2521628+ydshieh@users.noreply.github.com>

Co-authored-by: Yih-Dar <2521628+ydshieh@users.noreply.github.com>
---
 tests/models/jukebox/test_modeling_jukebox.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/models/jukebox/test_modeling_jukebox.py b/tests/models/jukebox/test_modeling_jukebox.py
index af5f946b9f..e77c8cb2eb 100644
--- a/tests/models/jukebox/test_modeling_jukebox.py
+++ b/tests/models/jukebox/test_modeling_jukebox.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import unittest
+from unittest import skip
 
 from transformers import is_torch_available
 from transformers.testing_utils import require_torch, slow
@@ -311,6 +312,7 @@ class Jukebox5bModelTester(unittest.TestCase):
         torch.testing.assert_allclose(zs[2][0], torch.tensor(self.EXPECTED_OUTPUT_0))
 
     @slow
+    @skip("Not enough GPU memory on CI runners")
     def test_slow_sampling(self):
         model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
         labels = [i.cuda() for i in self.prepare_inputs(self.model_id)]