From 56da20483701a8da2adb199631ef112f51dd270f Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Thu, 14 Sep 2023 22:15:36 -0400
Subject: [PATCH] use torch.cuda.current_device() instead of local_rank

---
 src/axolotl/utils/config.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/axolotl/utils/config.py b/src/axolotl/utils/config.py
index 90ed409b9c..a31f34b73e 100644
--- a/src/axolotl/utils/config.py
+++ b/src/axolotl/utils/config.py
@@ -29,7 +29,7 @@ def get_device():
         cfg.device_map = "auto"
     else:
         if cfg.device.startswith("cuda"):
-            cfg.device_map = {"": cfg.local_rank}
+            cfg.device_map = {"": torch.cuda.current_device()}
         else:
             cfg.device_map = {"": cfg.device}