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

[ColoTensor] improves init functions. (#1150)

Browse Source
This commit is contained in:
Jiarui Fang
2022-06-21 18:28:38 +08:00
committed by GitHub
parent 8106d7b8c7
commit 8cdce0399c
5 changed files with 103 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/tensor/colo_parameter.py
View File

@@ -35,7 +35,7 @@ class ColoParameter(ColoTensor, torch.nn.Parameter):
data: Optional[torch.Tensor] = None,
requires_grad: bool = True,
spec: TensorSpec = TensorSpec(distspec.replicate())) -> None:
self._spec = copy(spec)
self._tensor_spec = copy(spec)
self._type = TensorType.MODEL
self._graph_node = None