ros/LLaMA-Factory
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

tiny fix

Browse Source 
Former-commit-id: 37aa7099dff2a9a7b52e259dac92de41ce606946
This commit is contained in:
hiyouga
2023-12-01 15:58:50 +08:00
parent ec099b0586
commit 4738d002c7
2 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
src/llmtuner/extras/misc.py
View File

@@ -69,11 +69,12 @@ def count_parameters(model: torch.nn.Module) -> Tuple[int, int]:
def get_current_device() -> str:
import accelerate
local_rank = int(os.environ.get('LOCAL_RANK', '0'))
if accelerate.utils.is_xpu_available():
return "xpu:{}".format(local_rank)
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")
else:
return local_rank if torch.cuda.is_available() else "cpu"
return "cpu"
def get_logits_processor() -> "LogitsProcessorList":