Unverified Commit e34d9aa6 authored by Alex McKinney's avatar Alex McKinney Committed by GitHub
Browse files

Replaces `DIFFUSERS_TEST_DEVICE` backend list with trying device (#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.
parent 8d30d257
...@@ -44,13 +44,13 @@ if is_torch_available(): ...@@ -44,13 +44,13 @@ if is_torch_available():
if "DIFFUSERS_TEST_DEVICE" in os.environ: if "DIFFUSERS_TEST_DEVICE" in os.environ:
torch_device = os.environ["DIFFUSERS_TEST_DEVICE"] torch_device = os.environ["DIFFUSERS_TEST_DEVICE"]
try:
available_backends = ["cuda", "cpu", "mps"] # try creating device to see if provided device is valid
if torch_device not in available_backends: _ = torch.device(torch_device)
raise ValueError( except RuntimeError as e:
f"unknown torch backend for diffusers tests: {torch_device}. Available backends are:" raise RuntimeError(
f" {available_backends}" f"Unknown testing device specified by environment variable `DIFFUSERS_TEST_DEVICE`: {torch_device}"
) ) from e
logger.info(f"torch_device overrode to {torch_device}") logger.info(f"torch_device overrode to {torch_device}")
else: else:
torch_device = "cuda" if torch.cuda.is_available() else "cpu" torch_device = "cuda" if torch.cuda.is_available() else "cpu"
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment