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

[NFC] polish code style (#3268) 

Co-authored-by: Yuanchen Xu <yuanchen.xu00@gmail.com>
pull/3313/head
Yuanchen 2023-03-28 10:30:30 +08:00 committed by binmakeswell
parent c4a226b729
commit d58fa705b2
1 changed files with 5 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
colossalai/nn/_ops/_utils.py
View File

@ -1,12 +1,11 @@
import torch
from typing import Union, Optional, List
from colossalai.tensor import ColoTensor
from typing import List, Optional, Union
import torch
import torch.distributed as dist
from colossalai.global_variables import tensor_parallel_env as env
from colossalai.global_variables import tensor_parallel_env as env
from colossalai.nn.layer.utils import divide
from colossalai.tensor import ProcessGroup, ColoTensorSpec
from colossalai.tensor import ColoTensor, ColoTensorSpec, ProcessGroup
GeneralTensor = Union[ColoTensor, torch.Tensor]
Number = Union[int, float]
@ -135,7 +134,7 @@ class _ReduceInput(torch.autograd.Function):
class _SplitForwardGatherBackward(torch.autograd.Function):
"""
Split the input and keep only the corresponding chuck to the rank.
Args:
input_: input matrix.
process_group: parallel mode.