Commit d58fa705 authored by Yuanchen's avatar Yuanchen Committed by binmakeswell
Browse files

[NFC] polish code style (#3268)


Co-authored-by: default avatarYuanchen Xu <yuanchen.xu00@gmail.com>
parent c4a226b7
import torch from typing import List, Optional, Union
from typing import Union, Optional, List
from colossalai.tensor import ColoTensor
import torch import torch
import torch.distributed as dist 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.nn.layer.utils import divide
from colossalai.tensor import ProcessGroup, ColoTensorSpec from colossalai.tensor import ColoTensor, ColoTensorSpec, ProcessGroup
GeneralTensor = Union[ColoTensor, torch.Tensor] GeneralTensor = Union[ColoTensor, torch.Tensor]
Number = Union[int, float] Number = Union[int, float]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment