fix flaky cpu test

b5c684f0 · Patrick von Platen · da8e87e2 · b5c684f0
Commit b5c684f0 authored Jul 15, 2022 by Patrick von Platen
Show whitespace changes
Inline Side-by-side

Showing with 6 additions and 2 deletions

tests/test_modeling_utils.py tests/test_modeling_utils.py +6 -2

No files found.
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -1046,6 +1046,10 @@ class PipelineTesterMixin(unittest.TestCase):
        torch.manual_seed(0)
        image = sde_ve(num_inference_steps=2)
+        if model.device.type == "cpu":
+            expected_image_sum = 3384805888.0
+            expected_image_mean = 1076.00085
+        else:
            expected_image_sum = 3382849024.0
            expected_image_mean = 1075.3788