import torch


def encode_one_hot_tensor(labels):
    """ Takes a torch tensor of integers and encodes it into a one-hot tensor.
        Let K be the number of labels

        @param labels: a [T x H x W] torch tensor with values in {0, ..., K-1}

        @return: a [T x K x H x W] torch tensor of 0's and 1's
    """
    T, H, W = labels.shape
    K = int(torch.max(labels).item() + 1)
    
    # Encode the one hot tensor
    one_hot_tensor = torch.zeros((T, K, H, W), device=labels.device)
    one_hot_tensor.scatter_(1, labels.long().unsqueeze(1), 1)

    return one_hot_tensor