[Enhance] support to add custom settings to param_group (#283)

2022-06-09 20:11:19 +08:00 · 2022-06-09 20:11:19 +08:00 · 8b3675a2aa
parent a9afdad7a8
commit 8b3675a2aa
1 changed files with 3 additions and 0 deletions
--- a/mmengine/optim/optimizer/default_constructor.py
+++ b/mmengine/optim/optimizer/default_constructor.py
@ -216,6 +216,9 @@ class DefaultOptimWrapperConstructor:
                    if self.base_wd is not None:
                        decay_mult = custom_keys[key].get('decay_mult', 1.)
                        param_group['weight_decay'] = self.base_wd * decay_mult
+                    # add custom settings to param_group
+                    for k, v in custom_keys[key].items():
+                        param_group[k] = v
                    break

            if not is_custom: