Skip to content

GitLab

Commit b5c684f0 authored by Patrick von Platen's avatar Patrick von Platen
Browse files

fix flaky cpu test

parent da8e87e2
Hide whitespace changes
Inline Side-by-side
Showing with 6 additions and 2 deletions
tests/test_modeling_utils.py
View file @ b5c684f0
...@@ -1046,8 +1046,12 @@ class PipelineTesterMixin(unittest.TestCase): ...@@ -1046,8 +1046,12 @@ class PipelineTesterMixin(unittest.TestCase):
torch.manual_seed(0) torch.manual_seed(0)
image = sde_ve(num_inference_steps=2) image = sde_ve(num_inference_steps=2)
expected_image_sum = 3382849024.0 if model.device.type == "cpu":
expected_image_mean = 1075.3788 expected_image_sum = 3384805888.0
expected_image_mean = 1076.00085
else:
expected_image_sum = 3382849024.0
expected_image_mean = 1075.3788
assert (image.abs().sum() - expected_image_sum).abs().cpu().item() < 1e-2 assert (image.abs().sum() - expected_image_sum).abs().cpu().item() < 1e-2
assert (image.abs().mean() - expected_image_mean).abs().cpu().item() < 1e-4 assert (image.abs().mean() - expected_image_mean).abs().cpu().item() < 1e-4
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment