From 2ee3215949d8f2d3141c2340d8e4d24ec94b2384 Mon Sep 17 00:00:00 2001 From: Fanli Lin Date: Wed, 4 Sep 2024 10:34:03 +0800 Subject: [PATCH] [tests] make 2 tests device-agnostic (#9347) * enabel on xpu * fix style --- tests/models/autoencoders/test_models_vae.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/models/autoencoders/test_models_vae.py b/tests/models/autoencoders/test_models_vae.py index 5a0d62e7c9..1b1c9b3521 100644 --- a/tests/models/autoencoders/test_models_vae.py +++ b/tests/models/autoencoders/test_models_vae.py @@ -1036,9 +1036,9 @@ class ConsistencyDecoderVAEIntegrationTests(unittest.TestCase): "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/img2img/sketch-mountains-input.jpg" ).resize((256, 256)) - image = torch.from_numpy(np.array(image).transpose(2, 0, 1).astype(np.float32) / 127.5 - 1)[ - None, :, :, : - ].cuda() + image = torch.from_numpy(np.array(image).transpose(2, 0, 1).astype(np.float32) / 127.5 - 1)[None, :, :, :].to( + torch_device + ) latent = vae.encode(image).latent_dist.mean @@ -1079,7 +1079,7 @@ class ConsistencyDecoderVAEIntegrationTests(unittest.TestCase): image = ( torch.from_numpy(np.array(image).transpose(2, 0, 1).astype(np.float32) / 127.5 - 1)[None, :, :, :] .half() - .cuda() + .to(torch_device) ) latent = vae.encode(image).latent_dist.mean