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

fix device map

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

9
scripts/finetune.py
View File

@@ -47,10 +47,11 @@ 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":
cfg.device_map = {"": cfg.local_rank} if cfg.device.startswith("cuda"):
else: cfg.device_map = {"": cfg.local_rank}
cfg.device_map = {"": cfg.device} else:
cfg.device_map = {"": cfg.device}
def get_multi_line_input() -> Optional[str]: def get_multi_line_input() -> Optional[str]: