Skip to content

GitLab

Unverified Commit 6f6eef74 authored by Krishna Penukonda's avatar Krishna Penukonda Committed by GitHub
Browse files

Fix Compatibility with Nvidia NGC Containers (#919) 

Check if MPS backend is registered before calling is_available()
parent 8be48507
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 1 deletion
src/diffusers/utils/testing_utils.py
View file @ 6f6eef74
...@@ -29,7 +29,9 @@ if is_torch_available(): ...@@ -29,7 +29,9 @@ if is_torch_available():
) )
if is_torch_higher_equal_than_1_12: if is_torch_higher_equal_than_1_12:
torch_device = "mps" if torch.backends.mps.is_available() else torch_device # Some builds of torch 1.12 don't have the mps backend registered. See #892 for more details
mps_backend_registered = hasattr(torch.backends, "mps")
torch_device = "mps" if (mps_backend_registered and torch.backends.mps.is_available()) else torch_device
def get_tests_dir(append_path=None): def get_tests_dir(append_path=None):
......
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