[Kandinsky tests] add is_flaky to test_model_cpu_offload_forward_pa…

…ss (huggingface#6762) * add is_flaky to test_model_cpu_offload_forward_pass * style * update --------- Co-authored-by: Dhruv Nair <[email protected]>
stanleyjacob · Jan 31, 2024 · 0fc62d1 · 0fc62d1
1 parent f4d3f91
commit 0fc62d1
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py b/tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py
@@ -31,6 +31,7 @@
 from diffusers.utils.testing_utils import (
     enable_full_determinism,
     floats_tensor,
+    is_flaky,
     load_image,
     load_numpy,
     require_torch_gpu,
@@ -244,8 +245,9 @@ def test_inference_batch_single_identical(self):
     def test_float16_inference(self):
         super().test_float16_inference(expected_max_diff=5e-1)
 
+    @is_flaky()
     def test_model_cpu_offload_forward_pass(self):
-        super().test_inference_batch_single_identical(expected_max_diff=5e-4)
+        super().test_inference_batch_single_identical(expected_max_diff=8e-4)
 
     def test_save_load_optional_components(self):
         super().test_save_load_optional_components(expected_max_difference=5e-4)