update

pytorch · Nov 21, 2024 · bfef1a5 · bfef1a5
1 parent 2decb94
commit bfef1a5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/torchbenchmark/util/experiment/instantiator.py b/torchbenchmark/util/experiment/instantiator.py
@@ -122,7 +122,8 @@ def list_devices() -> List[str]:
     devices = ["cpu"]
     import torch
 
-    if device_type := torch._utils._get_available_device_type():
+    device_type = torch._C._get_accelerator().type
+    if device_type != "cpu":
         devices.append(device_type)
     return devices