From e5ab6be72ea3903e5b973aa7b31c9b3708ade3f2 Mon Sep 17 00:00:00 2001 From: jim <459633157@qq.com> Date: Mon, 10 Oct 2022 10:13:31 +0800 Subject: [PATCH] [hotfix[ fix colotensor.type() raise NotImplementedError (#1682) --- colossalai/nn/_ops/element_wise.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/colossalai/nn/_ops/element_wise.py b/colossalai/nn/_ops/element_wise.py index c3c1421e7..462670e72 100644 --- a/colossalai/nn/_ops/element_wise.py +++ b/colossalai/nn/_ops/element_wise.py @@ -18,6 +18,8 @@ def register_elementwise_op(op): output = op(input_tensor, *args, **kwargs) if isinstance(input_tensor, ColoTensor): + if isinstance(output, str): + return output if not isinstance(output, torch.Tensor): raise NotImplementedError return ColoTensor.from_torch_tensor(output,