Spaces:

Aatricks
/

LightDiffusion-Next

Running on Zero

Aatricks commited on Feb 4

Commit

c8b1d9f

1 Parent(s): 7cb9589

Update modules/Device/Device.py

Files changed (1) hide show

modules/Device/Device.py CHANGED Viewed

@@ -89,9 +89,9 @@ def get_torch_device() -> torch.device:
         if is_intel_xpu():
             return torch.device("xpu", torch.xpu.current_device())
         else:
-            try:
                 return torch.device(torch.cuda.current_device())
-            except AssertionError:
                 return torch.device("cpu")

         if is_intel_xpu():
             return torch.device("xpu", torch.xpu.current_device())
         else:
+            if is_nvidia():
                 return torch.device(torch.cuda.current_device())
+            else:
                 return torch.device("cpu")