github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-06 11:32:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

reorgnize colotensor directory (#1062)

Browse Source 
* reorgnize colotensor directory

* polish code
This commit is contained in:
Jiarui Fang
2022-06-03 18:04:22 +08:00
committed by GitHub
parent 3d10be33bd
commit a00644079e
25 changed files with 130 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/tensor/colo_parameter.py
View File

@@ -1,9 +1,9 @@
from .colo_tensor import ColoTensor
from .const import TensorType
from colossalai.tensor.colo_tensor import ColoTensor
from colossalai.tensor.const import TensorType
import torch
from colossalai.tensor import TensorSpec, distspec
from copy import copy
from .param_op_hook import _ParamOpHookWrapper, PreFwdPostBwd, PostFwdPreBwd
from colossalai.tensor.param_op_hook import _ParamOpHookWrapper, PreFwdPostBwd, PostFwdPreBwd
from typing import Optional