import torch


def one_hot(input_data:torch.Tensor, dim:int):
    """ Turns input_data of shape `H, W` with integer entries into a tensor of shape `H, C, W` where
    `C` is the one-hot encoding dimension. """
    res = []
    n_channels = input_data.size(1)
    offset = input_data.min()
    length = dim
    for channel_idx in range(n_channels):
        channel_one_hot = []
        channel = input_data[:, channel_idx]
        for entry in channel:
            one_hot_x = torch.zeros(length)
            one_hot_x[entry+offset] = 1
            channel_one_hot.append(one_hot_x)
        channel_one_hot = torch.cat(channel_one_hot)
        channel_one_hot = channel_one_hot.reshape(-1, length)
        res.append(channel_one_hot.unsqueeze(2))
    res = torch.cat(res, dim=2)
    return res