From 2f2803a443af7a9803823febac937006115a703a Mon Sep 17 00:00:00 2001 From: Ziheng Qin <37519855+henryqin1997@users.noreply.github.com> Date: Mon, 16 May 2022 09:52:24 -0600 Subject: [PATCH] [NFC] polish colossalai/nn/layer/utils/common.py code style --- colossalai/nn/layer/utils/common.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/colossalai/nn/layer/utils/common.py b/colossalai/nn/layer/utils/common.py index a112f2d95e56..f2297304fdc9 100644 --- a/colossalai/nn/layer/utils/common.py +++ b/colossalai/nn/layer/utils/common.py @@ -13,7 +13,8 @@ class CheckpointModule(nn.Module): - def __init__(self, checkpoint: bool = True, offload : bool = False): + + def __init__(self, checkpoint: bool = True, offload: bool = False): super().__init__() self.checkpoint = checkpoint self._use_checkpoint = checkpoint @@ -78,6 +79,7 @@ def get_tensor_parallel_mode(): def _ntuple(n): + def parse(x): if isinstance(x, collections.abc.Iterable): return x