From f1d7228a7472449df7159e6d69ca930256d74ce8 Mon Sep 17 00:00:00 2001 From: hiyouga Date: Fri, 1 Dec 2023 22:55:41 +0800 Subject: [PATCH] fix #1703 Former-commit-id: eee2e9abf6df345c5471e8ca7639293543ba720c --- src/llmtuner/extras/misc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llmtuner/extras/misc.py b/src/llmtuner/extras/misc.py index 33efb7d2..e1ae7d9f 100644 --- a/src/llmtuner/extras/misc.py +++ b/src/llmtuner/extras/misc.py @@ -73,7 +73,7 @@ def get_current_device() -> str: if accelerate.utils.is_xpu_available(): return "xpu:{}".format(os.environ.get("LOCAL_RANK", "0")) elif accelerate.utils.is_npu_available() or torch.cuda.is_available(): - return os.environ.get("LOCAL_RANK", "0") + return "cuda:{}".format(os.environ.get("LOCAL_RANK", "0")) else: return "cpu"