import torch


def one_hot(batched_tensor, num_class):
    """
    Convert a 1D batched tensor to one hot encoded tensors.

    Parameters
    ----------
    batched_tensor : torch.tensor
        Batched tensor of shape batch_size. The tensor type should be integers
        corresponding to the predicted class.
    num_class : int
        Number of classes.

    Returns
    -------
    one_hot : torch.tensor
        One hot encoded tensors.
    """
    one_hot = torch.zeros((batched_tensor.size()[0], num_class))
    one_hot[torch.arange(batched_tensor.size()[0]), batched_tensor.long()] = 1
    return one_hot