Fix vae tests for cpu and gpu (#480)

Author: kashif

tests/test_models_vae.py

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