def transform_to_ch_axis(x: _torch.Tensor, ch_axis: int):
    """
    Transform the given input tensor x so that the specified channel axis dimension
    becomes the outermost tensor dimension by permuting the tensor such that
    dimension 0 and the channel axis dimension are swapped. Then flatten all the
    dimensions from index 1 onwards in the permuted tensor and return the resulting tensor.

    The input tensor x is expected to be detached.
    """

    assert x.ndim >= 2

    assert ch_axis >= 0 and ch_axis < len(x.shape)

    new_axis_list = [i for i in range(len(x.shape))]
    new_axis_list[ch_axis] = 0
    new_axis_list[0] = ch_axis
    y = x.permute(new_axis_list)
    y = _torch.flatten(y, start_dim=1)

    return y

    Transform the given input tensor x so that the specified channel axis dimension
    becomes the outermost tensor dimension by permuting the tensor such that
    dimension 0 and the channel axis dimension are swapped. Then flatten all the
    dimensions from index 1 onwards in the permuted tensor and return the resulting tensor.

    The input tensor x is expected to be detached.
    
coremltools/optimize/torch/_utils/torch_utils.py