From 851474632d3836a9c1c8d543fdcc2dbe04e494de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BF=97=E6=B5=A9?= Date: Tue, 24 Sep 2024 23:26:34 +0800 Subject: [PATCH] add extract_token binary --- funasr/bin/extract_token.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/funasr/bin/extract_token.py b/funasr/bin/extract_token.py index 2e04d3a6f..2ddaf5299 100644 --- a/funasr/bin/extract_token.py +++ b/funasr/bin/extract_token.py @@ -23,7 +23,8 @@ def main_hydra(cfg: DictConfig): machine_rank = int(kwargs.get("dist_rank", 0)) world_size = int(kwargs.get("world_size", 1)) local_rank = int(os.environ["LOCAL_RANK"]) - node_world_size = int(os.environ["WORLD_SIZE"]) + node_world_size = int(kwargs.get("num_gpus", 1)) + logging.info(f"machine_rank: {machine_rank}, world_size: {world_size}, local_rank: {local_rank}, node_world_size: {node_world_size}") dist_rank = machine_rank * node_world_size + local_rank logging.basicConfig( level='INFO',