def flatten(tensor):
    """Flattens a given tensor such that the channel axis is first.
    The shapes are transformed as follows:
    (N, C, D, H, W) -> (C, N * D * H * W)
    """
    # new axis order
    axis_order = (1, 0) + tuple(range(2, len(tensor.shape)))
    # Transpose: (N, C, D, H, W) -> (C, N, D, H, W)
    transposed = paddle.transpose(tensor, perm=axis_order)
    # Flatten: (C, N, D, H, W) -> (C, N * D * H * W)
    return paddle.flatten(transposed, start_axis=1, stop_axis=-1)
Flattens a given tensor such that the channel axis is first.
    The shapes are transformed as follows:
    (N, C, D, H, W) -> (C, N * D * H * W)
    
contrib/MedicalSeg/medicalseg/models/losses/loss_utils.py