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

[NFC] polish colossalai/auto_parallel/tensor_shard/utils/factory.py code style (#2349)

pull/2359/head
LuGY 2023-01-05 21:17:42 +08:00 committed by GitHub
parent 28e2d16794
commit e11a005c02
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/auto_parallel/tensor_shard/utils/factory.py
View File

@ -4,10 +4,11 @@ from functools import reduce
from typing import Dict, List, Optional, Union from typing import Dict, List, Optional, Union
import torch import torch
from torch.fx.node import Node
from colossalai.device.device_mesh import DeviceMesh from colossalai.device.device_mesh import DeviceMesh
from colossalai.tensor.shape_consistency import ShapeConsistencyManager from colossalai.tensor.shape_consistency import ShapeConsistencyManager
from colossalai.tensor.sharding_spec import ShardingSpec from colossalai.tensor.sharding_spec import ShardingSpec
from torch.fx.node import Node
from ..constants import INFINITY_COST from ..constants import INFINITY_COST