import torch


def _onehot_encode(predictions: torch.Tensor,
        labels: torch.Tensor,
        n_classes: int):
    """
    One-hot-encode the predictions and the labels.
    :param predictions: BxHxW int64 tensor containing the class predictions
    :param labels: BxHxW int64 tensor containing the correct class labels
    :param n_classes: number of classes present
    :return:
    """
    print('n classes: ', n_classes)
    pred_tensor = torch.nn.functional.one_hot(predictions, num_classes=n_classes).permute(0, 3, 1, 2)
    onehot_gt_tensor = torch.nn.functional.one_hot(labels, num_classes=n_classes).permute(0, 3, 1, 2)
    return pred_tensor, onehot_gt_tensor