From 52d145a34249028b0c918456f2f88f75ca27201a Mon Sep 17 00:00:00 2001
From: Super Daniel <78588128+super-dainiu@users.noreply.github.com>
Date: Tue, 12 Jul 2022 18:56:34 +0800
Subject: [PATCH] [NFC] polish colossalai/nn/lr_scheduler/onecycle.py code
 style (#1269)

---
 colossalai/nn/lr_scheduler/onecycle.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/colossalai/nn/lr_scheduler/onecycle.py b/colossalai/nn/lr_scheduler/onecycle.py
index b319e7bb6..8007fd360 100644
--- a/colossalai/nn/lr_scheduler/onecycle.py
+++ b/colossalai/nn/lr_scheduler/onecycle.py
@@ -68,7 +68,9 @@ class OneCycleLR(_OneCycleLR):
         https://arxiv.org/abs/1708.07120
     """
 
-    def __init__(self, optimizer, total_steps: int,
+    def __init__(self,
+                 optimizer,
+                 total_steps: int,
                  pct_start=0.3,
                  anneal_strategy='cos',
                  cycle_momentum=True,
@@ -76,9 +78,12 @@ class OneCycleLR(_OneCycleLR):
                  max_momentum=0.95,
                  div_factor=25.0,
                  final_div_factor=10000.0,
-                 last_epoch=-1, **kwargs):
+                 last_epoch=-1,
+                 **kwargs):
         max_lrs = list(map(lambda group: group['lr'], optimizer.param_groups))
-        super().__init__(optimizer, max_lrs, total_steps=total_steps,
+        super().__init__(optimizer,
+                         max_lrs,
+                         total_steps=total_steps,
                          pct_start=pct_start,
                          anneal_strategy=anneal_strategy,
                          cycle_momentum=cycle_momentum,