From 7eb87f516dfd242684b24c743f550ad9c0caf02c Mon Sep 17 00:00:00 2001
From: Liang Bowen <82053485+number1roy@users.noreply.github.com>
Date: Wed, 9 Mar 2022 17:34:43 +0800
Subject: [PATCH] flake8 style (#352)

---
 colossalai/nn/layer/utils/common.py             | 2 +-
 colossalai/nn/layer/vanilla/layers.py           | 2 +-
 colossalai/nn/layer/wrapper/pipeline_wrapper.py | 6 ++++--
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/colossalai/nn/layer/utils/common.py b/colossalai/nn/layer/utils/common.py
index c1d88d2fc..2ec7d19dd 100644
--- a/colossalai/nn/layer/utils/common.py
+++ b/colossalai/nn/layer/utils/common.py
@@ -38,7 +38,7 @@ class CheckpointModule(nn.Module):
 
 def divide(numerator, denominator):
     """Only allow exact division
-    
+
     :param numerator: Numerator of the division
     :param denominator: Denominator of the division
     """
diff --git a/colossalai/nn/layer/vanilla/layers.py b/colossalai/nn/layer/vanilla/layers.py
index e5c9fd074..2f2c441e1 100644
--- a/colossalai/nn/layer/vanilla/layers.py
+++ b/colossalai/nn/layer/vanilla/layers.py
@@ -101,7 +101,7 @@ class WrappedDropPath(nn.Module):
 
 @LAYERS.register_module
 class VanillaPatchEmbedding(nn.Module):
-    """ 
+    """
     2D Image to Patch Embedding
 
     :param img_size: image size
diff --git a/colossalai/nn/layer/wrapper/pipeline_wrapper.py b/colossalai/nn/layer/wrapper/pipeline_wrapper.py
index dd422a75f..20813dc89 100644
--- a/colossalai/nn/layer/wrapper/pipeline_wrapper.py
+++ b/colossalai/nn/layer/wrapper/pipeline_wrapper.py
@@ -33,14 +33,16 @@ class PipelineSharedModuleWrapper:
                     self.ranks_in_group = sub_ranks
 
     def register_module(self, module: nn.Module):
-        assert self.ranks_in_group is not None, f'Rank {gpc.get_local_rank(ParallelMode.PIPELINE)} is not in pipeline_ranks {self.pipeline_ranks}'
+        assert self.ranks_in_group is not None,\
+            f'Rank {gpc.get_local_rank(ParallelMode.PIPELINE)} is not in pipeline_ranks {self.pipeline_ranks}'
         src = self.ranks_in_group[self.pipeline_ranks[0]]
         for p in module.parameters():
             setattr(p, 'pipeline_shared_module_pg', self.group)
             dist.broadcast(p, src, group=self.group)
 
     def register_parameter(self, param: nn.Parameter):
-        assert self.ranks_in_group is not None, f'Rank {gpc.get_local_rank(ParallelMode.PIPELINE)} is not in pipeline_ranks {self.pipeline_ranks}'
+        assert self.ranks_in_group is not None,\
+            f'Rank {gpc.get_local_rank(ParallelMode.PIPELINE)} is not in pipeline_ranks {self.pipeline_ranks}'
         src = self.ranks_in_group[self.pipeline_ranks[0]]
         setattr(param, 'pipeline_shared_module_pg', self.group)
         dist.broadcast(param, src, group=self.group)