github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[NFC] polish colossalai/nn/_ops/layernorm.py code style (#1555)

pull/1550/head
yuxuan-lou 2022-09-08 15:15:52 +08:00 committed by Frank Lee
parent 8edb777cc2
commit 413f9c19f4
1 changed files with 3 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
colossalai/nn/_ops/layernorm.py
View File

@ -19,9 +19,7 @@ def colo_layernorm(
input_tensor = input_tensor.redistribute(ReplicaSpec()) input_tensor = input_tensor.redistribute(ReplicaSpec())
output = F.layer_norm(input_tensor, normalized_shape, weight=weight, bias=bias, eps=eps) output = F.layer_norm(input_tensor, normalized_shape, weight=weight, bias=bias, eps=eps)
output = ColoTensor.from_torch_tensor( output = ColoTensor.from_torch_tensor(tensor=output,
tensor=output, spec=ColoTensorSpec(pg=input_tensor.get_process_group(),
spec=ColoTensorSpec( dist_attr=input_tensor.dist_spec))
pg=input_tensor.get_process_group(),
dist_attr=input_tensor.dist_spec))
return output return output