github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-08 04:24:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Tensor] init a tp network training unittest (#849)

Browse Source
This commit is contained in:
Jiarui Fang
2022-04-24 16:43:44 +08:00
committed by GitHub
parent 0dea140760
commit 62f059251b
6 changed files with 113 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/utils/model/colo_init_context.py
View File

@@ -26,4 +26,4 @@ class ColoInitContext(InsertPostInitMethodToModuleSubClasses):
save_torch_payload = True if not self._lazy_memory_allocate else False
for name, param in name_list:
delattr(module, name)
setattr(module, name, ColoTensor.init_from_torch_tensor(tensor=param.data, save_payload=save_torch_payload))
setattr(module, name, ColoTensor.init_from_torch_tensor(tensor=param, save_payload=save_torch_payload))