fix #1047: process_device will return input when device is working. (#…

…1143)
sbi-dev · Apr 12, 2024 · 0b5f931 · 0b5f931
1 parent 81ce2c9
commit 0b5f931
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/tests/torchutils_test.py b/tests/torchutils_test.py
@@ -215,10 +215,8 @@ def test_process_device(device_input: str) -> None:
             elif torch.backends.mps.is_available():
                 assert device_output == "mps:0"
 
-        if device_input == "cuda" and torch.cuda.is_available():
-            assert device_output == "cuda:0"
-        if device_input == "cuda:0" and torch.cuda.is_available():
-            assert device_output == "cuda:0"
+        if device_input.startswith("cuda") and torch.cuda.is_available():
+            assert device_output == device_input
         if device_input == "mps" and torch.backends.mps.is_available():
             assert device_output == "mps"