import torch


def confusion_matrix(input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """Calculate confusion matrix

    :param input: output of network, expected to be `BxCx(OPTIONAL DIMENSIONS)`
    :param target: target, expected to be `Bx(OPTIONAL DIMENSIONS)`
    :return: confusion matrix in long tensor of `CxC`
    """
    num_classes = input.size(1)
    classes = torch.arange(num_classes, device=input.device)
    pred = input.argmax(dim=1).view(-1, 1)
    target = target.view(-1, 1)
    return ((pred == classes).t() @ (target == classes)).long()