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] rename APIs and add output_replicate to ComputeSpec (#1168)

Browse Source
This commit is contained in:
Jiarui Fang
2022-06-24 13:08:54 +08:00
committed by GitHub
parent f4ef224358
commit 4b9bba8116
23 changed files with 116 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/tensor/colo_parameter.py
View File

@@ -82,7 +82,7 @@ class ColoParameter(ColoTensor, torch.nn.Parameter):
else:
with torch._C.DisableTorchFunction():
data = self.data.clone()
tensor = ColoParameter(data, self.requires_grad, spec=copy(self.spec))
tensor = ColoParameter(data, self.requires_grad, spec=copy(self.tensor_spec))
memo[id(self)] = tensor
return tensor