tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Merge pull request #148 from OpenAccess-AI-Collective/fix-device-load

Browse Source
This commit is contained in:
Wing Lian
2023-06-02 14:37:17 -04:00
committed by GitHub
parent 76a70fd739 74ebbf4371
commit 4ab0c8b201
1 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
scripts/finetune.py
View File

@@ -47,7 +47,8 @@ def choose_device(cfg):
return "cpu" return "cpu"
cfg.device = get_device() cfg.device = get_device()
if cfg.device == "cuda": if cfg.device_map != "auto":
if cfg.device.startswith("cuda"):
cfg.device_map = {"": cfg.local_rank} cfg.device_map = {"": cfg.local_rank}
else: else:
cfg.device_map = {"": cfg.device} cfg.device_map = {"": cfg.device}