enable find_unused_parameters flag for DDP

2023-11-21 19:26:42 +08:00
parent 112208ef74
commit 388974ab2a
3 changed files with 4 additions and 3 deletions
@@ -46,7 +46,7 @@ def run_model(cfgs, training):
        model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
    if cfgs['trainer_cfg']['fix_BN']:
        model.fix_BN()
-    model = get_ddp_module(model)
+    model = get_ddp_module(model, cfgs['trainer_cfg']['find_unused_parameters'])
    msg_mgr.log_info(params_count(model))
    msg_mgr.log_info("Model Initialization Finished!")