diff --git a/tests/models/idefics2/test_modeling_idefics2.py b/tests/models/idefics2/test_modeling_idefics2.py index d1a982dc595cb1..8d48acb9500d2f 100644 --- a/tests/models/idefics2/test_modeling_idefics2.py +++ b/tests/models/idefics2/test_modeling_idefics2.py @@ -34,6 +34,7 @@ require_flash_attn, require_torch, require_torch_gpu, + require_torch_multi_gpu, slow, torch_device, ) @@ -498,7 +499,7 @@ def tearDown(self): torch.cuda.empty_cache() @slow - @unittest.skip("Test hits OOM on CI - https://github.com/huggingface/transformers/issues/32288") + @require_torch_multi_gpu def test_integration_test(self): model = Idefics2ForConditionalGeneration.from_pretrained( "HuggingFaceM4/idefics2-8b-base",