typing

Author: Borda

src/lightning/pytorch/callbacks/lr_monitor.py

@@ -255,9 +255,8 @@ def _extract_momentum(self, param_group: dict[str, list[float]], name: str, use_
         if not self.log_momentum:
             return {}
 
-        momentum = param_group["betas"][0] if use_betas else param_group.get("momentum")
-        momentum = momentum if momentum is not None else 0
-        self.last_momentum_values[name] = momentum
+        momentum = param_group["betas"][0] if use_betas else param_group.get("momentum", 0)
+        self.last_momentum_values[name] = momentum  # type: ignore[assignment]
         return {name: momentum}
 
     def _extract_weight_decay(self, param_group: dict[str, Any], name: str) -> dict[str, Any]: