From da7e3994ad3408031113ef661650c625320d30cf Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Tue, 13 Sep 2022 19:14:20 +0200 Subject: [PATCH] Fix vae tests for cpu and gpu (#480) --- tests/test_models_vae.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/test_models_vae.py b/tests/test_models_vae.py index cc4f65c3..361eb618 100644 --- a/tests/test_models_vae.py +++ b/tests/test_models_vae.py @@ -104,7 +104,15 @@ class AutoencoderKLTests(ModelTesterMixin, unittest.TestCase): output_slice = output[0, -1, -3:, -3:].flatten().cpu() - # fmt: off - expected_output_slice = torch.tensor([-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]) - # fmt: on + # Since the VAE Gaussian prior's generator is seeded on the appropriate device, + # the expected output slices are not the same for CPU and GPU. + if torch_device in ("mps", "cpu"): + expected_output_slice = torch.tensor( + [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026] + ) + else: + expected_output_slice = torch.tensor( + [-0.2421, 0.4642, 0.2507, -0.0438, 0.0682, 0.3160, -0.2018, -0.0727, 0.2485] + ) + self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))