From 0cf8e8e91c9b8cfdba0e50deb609c69b08c9ac76 Mon Sep 17 00:00:00 2001
From: oahzxl <43881818+oahzxl@users.noreply.github.com>
Date: Tue, 12 Jul 2022 18:18:14 +0800
Subject: [PATCH] [NFC] polish <colossalai/nn/lr_scheduler/poly.py> code style
 (#1267)

---
 colossalai/nn/lr_scheduler/poly.py | 28 ++++++++++++++++++----------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/colossalai/nn/lr_scheduler/poly.py b/colossalai/nn/lr_scheduler/poly.py
index 5e4ee7150..16352bc51 100644
--- a/colossalai/nn/lr_scheduler/poly.py
+++ b/colossalai/nn/lr_scheduler/poly.py
@@ -17,7 +17,12 @@ class PolynomialLR(_LRScheduler):
             the schedule is started from the beginning or When last_epoch=-1, sets initial lr as lr.
     """
 
-    def __init__(self, optimizer, total_steps: int, end_lr: float = 0.0001, power: float = 1.0, last_epoch: int = -1,
+    def __init__(self,
+                 optimizer,
+                 total_steps: int,
+                 end_lr: float = 0.0001,
+                 power: float = 1.0,
+                 last_epoch: int = -1,
                  **kwargs):
         if end_lr < 0:
             raise ValueError(f'end_lr must >= 0, got {end_lr}')
@@ -30,11 +35,9 @@ class PolynomialLR(_LRScheduler):
         return self._get_closed_form_lr()
 
     def _get_closed_form_lr(self):
-        return [
-            (base_lr - self.end_lr) * ((1 - min(self.last_epoch, self.total_steps) /
-                                        self.total_steps) ** self.power) + self.end_lr
-            for base_lr in self.base_lrs
-        ]
+        return [(base_lr - self.end_lr) *
+                ((1 - min(self.last_epoch, self.total_steps) / self.total_steps)**self.power) + self.end_lr
+                for base_lr in self.base_lrs]
 
 
 @LR_SCHEDULERS.register_module
@@ -51,8 +54,13 @@ class PolynomialWarmupLR(WarmupScheduler):
             the schedule is started from the beginning or When last_epoch=-1, sets initial lr as lr.
     """
 
-    def __init__(self, optimizer, total_steps: int, warmup_steps: int = 0, end_lr: float = 0.0001, power: float = 1.0,
-                 last_epoch: int = -1, **kwargs):
-        base_scheduler = PolynomialLR(
-            optimizer, total_steps - warmup_steps, end_lr=end_lr, power=power)
+    def __init__(self,
+                 optimizer,
+                 total_steps: int,
+                 warmup_steps: int = 0,
+                 end_lr: float = 0.0001,
+                 power: float = 1.0,
+                 last_epoch: int = -1,
+                 **kwargs):
+        base_scheduler = PolynomialLR(optimizer, total_steps - warmup_steps, end_lr=end_lr, power=power)
         super().__init__(optimizer, warmup_steps, base_scheduler, last_epoch=last_epoch)