From 86ada491e01691c53fd72fa76b8c77294042f938 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B8=B8=E9=9B=81?= Date: Fri, 17 May 2024 12:38:24 +0800 Subject: [PATCH] deepspeed --- funasr/bin/train_ds.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/funasr/bin/train_ds.py b/funasr/bin/train_ds.py index b0931b016..03cb18479 100644 --- a/funasr/bin/train_ds.py +++ b/funasr/bin/train_ds.py @@ -130,8 +130,8 @@ def main(**kwargs): model = trainer.warp_model(model) - kwargs["device"] = next(model.parameters()).device - trainer.device = kwargs["device"] + kwargs["device"] = int(os.environ.get("LOCAL_RANK", 0)) + trainer.device = int(os.environ.get("LOCAL_RANK", 0)) model, optim, scheduler = trainer.warp_optim_scheduler(model, **kwargs)