Skip to content

Commit

Permalink
Fix vae tests for cpu and gpu (huggingface#480)
Browse files Browse the repository at this point in the history
  • Loading branch information
kashif authored Sep 13, 2022
1 parent 55f7ca3 commit da7e399
Showing 1 changed file with 11 additions and 3 deletions.
14 changes: 11 additions & 3 deletions tests/test_models_vae.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,15 @@ def test_output_pretrained(self):

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))

0 comments on commit da7e399

Please sign in to comment.