Skip to content

Commit

Permalink
[tests] make 2 tests device-agnostic (#9347)
Browse files Browse the repository at this point in the history
* enabel on xpu

* fix style
  • Loading branch information
faaany authored Sep 4, 2024
1 parent 8ecf499 commit 2ee3215
Showing 1 changed file with 4 additions and 4 deletions.
8 changes: 4 additions & 4 deletions tests/models/autoencoders/test_models_vae.py
Original file line number Diff line number Diff line change
Expand Up @@ -1036,9 +1036,9 @@ def test_encode_decode(self):
"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

Expand Down Expand Up @@ -1079,7 +1079,7 @@ def test_encode_decode_f16(self):
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
Expand Down

0 comments on commit 2ee3215

Please sign in to comment.