SchedulerCallback: save lr and momentum only for 0-th param group (#1028)

* SchedulerCallback: save lr and momentum only for 0-th param group

* Codestyle

Author: and-kul

catalyst/callbacks/scheduler.py

@@ -136,48 +136,32 @@ def _update_lr_and_momentum_in_metrics_dict(
         momentum_list: List[Union[float, None]],
     ):
         """Update learning rate and momentum in metrics_dict
+        (consider only 0-th param group)
 
         Args:
             metrics_dict: batch_metrics or epoch_metrics
             lr_list: lr for each param group
             momentum_list: momentum for each param group
 
         """
-        if len(lr_list) == 1:
-            lr = lr_list[0]
-            momentum = momentum_list[0]
+        # todo: consider saving lr and momentum for all param groups ?
+        lr = lr_list[0]
+        momentum = momentum_list[0]
+
+        lr_key = (
+            f"lr/{self.scheduler_key}"
+            if self.scheduler_key is not None
+            else "lr"
+        )
+        metrics_dict[lr_key] = lr
 
-            lr_key = (
-                f"lr/{self.scheduler_key}"
+        if momentum is not None:
+            momentum_key = (
+                f"momentum/{self.scheduler_key}"
                 if self.scheduler_key is not None
-                else "lr"
+                else "momentum"
             )
-            metrics_dict[lr_key] = lr
-
-            if momentum is not None:
-                momentum_key = (
-                    f"momentum/{self.scheduler_key}"
-                    if self.scheduler_key is not None
-                    else "momentum"
-                )
-                metrics_dict[momentum_key] = momentum
-
-        else:
-            for i, (lr, momentum) in enumerate(zip(lr_list, momentum_list)):
-                lr_key = (
-                    f"lr/{self.scheduler_key}/group_{i}"
-                    if self.scheduler_key is not None
-                    else f"lr/group_{i}"
-                )
-                metrics_dict[lr_key] = lr
-
-                if momentum is not None:
-                    momentum_key = (
-                        f"momentum/{self.scheduler_key}/group_{i}"
-                        if self.scheduler_key is not None
-                        else f"momentum/group_{i}"
-                    )
-                    metrics_dict[momentum_key] = momentum
+            metrics_dict[momentum_key] = momentum
 
     def step_batch(self, runner: "IRunner") -> None:
         """Perform scheduler step and update batch metrics