def METHOD_NAME(op):
    @colo_op_impl(op)
    def elementwise_op(input_tensor: GeneralTensor, *args, **kwargs):
        """
        Handles ``__torch_function__`` dispatch for the elementwise op such
        as ``torch.nn.functional.gelu`` or ``torch.nn.functional.relu``.
        This method computes on either a normal tensor or a sharded tensor.
        """
        if 'inplace' in kwargs:
            # TODO(jiaruifang) inplace will cause bugs
            input_tensor = input_tensor.clone()
            return op(input_tensor, *args, **kwargs)
        else:
            output = op(input_tensor, *args, **kwargs)
            # return output
            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,
                                                    spec=ColoTensorSpec(input_tensor.get_process_group(),
                                                                        dist_attr=input_tensor.dist_spec))