From 413f9c19f4f05bd714ed112903a94b1071327bab Mon Sep 17 00:00:00 2001 From: yuxuan-lou <83441848+yuxuan-lou@users.noreply.github.com> Date: Thu, 8 Sep 2022 15:15:52 +0800 Subject: [PATCH] [NFC] polish colossalai/nn/_ops/layernorm.py code style (#1555) --- colossalai/nn/_ops/layernorm.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/colossalai/nn/_ops/layernorm.py b/colossalai/nn/_ops/layernorm.py index e3eef9b18..2b761b84e 100644 --- a/colossalai/nn/_ops/layernorm.py +++ b/colossalai/nn/_ops/layernorm.py @@ -19,9 +19,7 @@ def colo_layernorm( input_tensor = input_tensor.redistribute(ReplicaSpec()) output = F.layer_norm(input_tensor, normalized_shape, weight=weight, bias=bias, eps=eps) - output = ColoTensor.from_torch_tensor( - tensor=output, - spec=ColoTensorSpec( - pg=input_tensor.get_process_group(), - dist_attr=input_tensor.dist_spec)) + output = ColoTensor.from_torch_tensor(tensor=output, + spec=ColoTensorSpec(pg=input_tensor.get_process_group(), + dist_attr=input_tensor.dist_spec)) return output