From 46a27fa5e6ad91544b7bc78240ffb07f687593ba Mon Sep 17 00:00:00 2001 From: Henry Date: Wed, 22 May 2024 20:17:36 +0800 Subject: [PATCH] fix rank --- .../troubleshooter/migrator/api_dump/universal_interface.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/troubleshooter/troubleshooter/migrator/api_dump/universal_interface.py b/troubleshooter/troubleshooter/migrator/api_dump/universal_interface.py index eb5f657..97c656b 100644 --- a/troubleshooter/troubleshooter/migrator/api_dump/universal_interface.py +++ b/troubleshooter/troubleshooter/migrator/api_dump/universal_interface.py @@ -117,7 +117,7 @@ def api_dump_init(net, output_path=os.path.join(os.getcwd(), "ts_api_dump"), *, elif "torch" in FRAMEWORK_TYPE and isinstance(net, torch.nn.Module): API_DUMP_FRAMEWORK_TYPE = "torch" pt_set_dump_path(output_path) - pt_register_hook(net, pt_acc_cmp_dump, compare_statedict=compare_statedict, rank=os.getenv("RANK_ID")) + pt_register_hook(net, pt_acc_cmp_dump, compare_statedict=compare_statedict, rank=os.getenv("RANK")) elif "mindspore" in FRAMEWORK_TYPE and isinstance(net, mindspore.nn.Cell): API_DUMP_FRAMEWORK_TYPE = "mindspore" ms_set_dump_path(output_path) -- Gitee