Fixed issue with cuda index

- Fixed warning with wandb
makgyver · Oct 4, 2024 · 1b6dc65 · 1b6dc65
1 parent 3b4ec98
commit 1b6dc65
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/fluke/__init__.py b/fluke/__init__.py
@@ -330,8 +330,8 @@ def set_device(self, device: str) -> torch.device:
             return GlobalSettings().auto_device()
 
         if device.startswith('cuda') and ":" in device:
-            idx = int(device.split(":")[1])
-            self._device = torch.device(device, idx)
+            idx = int(device.split(":")[0])
+            self._device = torch.device("cuda", idx)
         else:
             self._device = torch.device(device)
         return self._device

diff --git a/fluke/utils/log.py b/fluke/utils/log.py
@@ -24,8 +24,6 @@
 from ..comm import ChannelObserver, Message  # NOQA
 from . import ClientObserver, ServerObserver, get_class_from_str  # NOQA
 
-wandb.require("core")
-
 
 __all__ = [
     "Log",