github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-15 06:33:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
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