import torch


def make_one_hot(tensor, n=None):
    """
    Convert a label map into a one hot vector map.

    :param tensor: Tensor indicating class of each item (b, 1, h, w)
    :param n: Number of classes (default uses max value in tensor)
    :return: Tensor with one hot vectors, i.e. a 1 at the class index (b, n, h, w)
    """
    n, (b, _, h, w) = n or tensor.max(), tensor.shape
    onehot = torch.zeros((b, n, h, w), device=tensor.device, dtype=torch.long)
    onehot.scatter_(1, tensor, 1)
    return onehot