423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-12-15 19:30:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

tiny fix

Browse Source
This commit is contained in:
hiyouga
2023-12-01 15:58:50 +08:00
parent fbc6220692
commit e597d3c084
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":