github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[hotfix[ fix colotensor.type() raise NotImplementedError (#1682)

pull/1683/head
jim 2022-10-10 10:13:31 +08:00 committed by GitHub
parent 3b2a59b0ba
commit e5ab6be72e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/nn/_ops/element_wise.py
View File

@ -18,6 +18,8 @@ def register_elementwise_op(op):
output = op(input_tensor, *args, **kwargs) output = op(input_tensor, *args, **kwargs)
if isinstance(input_tensor, ColoTensor): if isinstance(input_tensor, ColoTensor):
if isinstance(output, str):
return output
if not isinstance(output, torch.Tensor): if not isinstance(output, torch.Tensor):
raise NotImplementedError raise NotImplementedError
return ColoTensor.from_torch_tensor(output, return ColoTensor.from_torch_tensor(output,