Fix Compatibility with Nvidia NGC Containers (#919)

Check if MPS backend is registered before calling is_available()

Fix Compatibility with Nvidia NGC Containers (#919)
Check if MPS backend is registered before calling is_available()
6f6eef74 · Krishna Penukonda · GitHub · 8be48507 · 6f6eef74
Unverified Commit 6f6eef74 authored Oct 21, 2022 by Krishna Penukonda Committed by GitHub Oct 20, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

src/diffusers/utils/testing_utils.py src/diffusers/utils/testing_utils.py +3 -1

No files found.
--- a/src/diffusers/utils/testing_utils.py
+++ b/src/diffusers/utils/testing_utils.py
@@ -29,7 +29,9 @@ if is_torch_available():
    )

    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):