import torch


def flatten(tensor: torch.Tensor) -> torch.Tensor:
    """
    Flattens a given tensor into a two-dimensional tensor. Class channel becomes the first dimension and
    other dimensions are squashed into one.

       3D: (Batch, Class, Depth, Height, Width) -> (C, B * D * H * W)\n
       2D: (B, C, H, W) -> (C, B * H * W)
    """

    num_classes = tensor.size()[1]
    # new dimension order
    axis_order = (1, 0) + tuple(range(2, tensor.dim()))
    # Transpose: (N, C, D, H, W) -> (C, N, D, H, W)
    transposed = tensor.permute(axis_order)
    # Flatten: (C, N, D, H, W) -> (C, N * D * H * W)
    return transposed.contiguous().view(num_classes, -1)