[NFC] polish colossalai/nn/layer/colossalai_layer/dropout.py code style (#1568)

This commit is contained in:
DouJS 2022-09-08 16:23:04 +08:00 committed by Frank Lee
parent c7d4932956
commit f586887a90

View File

@ -13,6 +13,7 @@ class Dropout(ColossalaiModule):
p (float, optional): probability of an element to be zeroed, defaults 0.5. p (float, optional): probability of an element to be zeroed, defaults 0.5.
inplace (bool, optional): whether to do dropout in-place, default to be False. inplace (bool, optional): whether to do dropout in-place, default to be False.
""" """
def __init__(self, p: float = 0.5, inplace: bool = False) -> None: def __init__(self, p: float = 0.5, inplace: bool = False) -> None:
tensor_parallel = get_tensor_parallel_mode() tensor_parallel = get_tensor_parallel_mode()
if tensor_parallel == "1d": if tensor_parallel == "1d":