Skip to content

Commit

Permalink
Replaces DIFFUSERS_TEST_DEVICE backend list with trying device (hug…
Browse files Browse the repository at this point in the history
…gingface#4673)

This is a better method than comparing against a list of supported backends as it allows for supporting any number of backends provided they are installed on the user's system.
This should have no effect on the behaviour of tests in Huggingface's CI workers.
See transformers#25506 where this approach has already been added.
  • Loading branch information
vvvm23 authored Aug 22, 2023
1 parent 8d30d25 commit e34d9aa
Showing 1 changed file with 7 additions and 7 deletions.
14 changes: 7 additions & 7 deletions src/diffusers/utils/testing_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,13 +44,13 @@

if "DIFFUSERS_TEST_DEVICE" in os.environ:
torch_device = os.environ["DIFFUSERS_TEST_DEVICE"]

available_backends = ["cuda", "cpu", "mps"]
if torch_device not in available_backends:
raise ValueError(
f"unknown torch backend for diffusers tests: {torch_device}. Available backends are:"
f" {available_backends}"
)
try:
# try creating device to see if provided device is valid
_ = torch.device(torch_device)
except RuntimeError as e:
raise RuntimeError(
f"Unknown testing device specified by environment variable `DIFFUSERS_TEST_DEVICE`: {torch_device}"
) from e
logger.info(f"torch_device overrode to {torch_device}")
else:
torch_device = "cuda" if torch.cuda.is_available() else "cpu"
Expand Down

0 comments on commit e34d9aa

Please sign in to comment.