import torch


def to_one_hot(x: torch.Tensor, num_vals) -> torch.Tensor:
    """
    Convert tensor to one-hot encoding
    """
    if type(x) is not torch.Tensor:
        x = torch.tensor(x)
    x = x.type(torch.long)
    x_one_hot = torch.zeros((x.shape[0], num_vals), dtype=torch.float32)
    x_one_hot = x_one_hot.scatter(1, x, 1.)
    return x_one_hot