Update __init__.py

makgyver · Oct 4, 2024 · 76bf37c · 76bf37c
1 parent a7a1aec
commit 76bf37c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/fluke/__init__.py b/fluke/__init__.py
@@ -330,7 +330,7 @@ def set_device(self, device: str) -> torch.device:
             return GlobalSettings().auto_device()
 
         if device.startswith('cuda') and ":" in device:
-            idx = int(device.split(":")[0])
+            idx = int(device.split(":")[1])
             self._device = torch.device("cuda", idx)
         else:
             self._device = torch.device(device)