import torch


def tensor_extend_zero(x, dim=0):
    """Extending tensor with zeros along an axis.

    Args:
        x (Tensor): tensor to be extended
        dim (int): the axis to extend zeros
    Returns:
        Tensor: the extended tensor. Its shape is
            ``(*x.shape[:dim], x.shape[dim]+1, *x.shape[dim+1:])``
    """
    zero_shape = list(x.shape)
    zero_shape[dim] = 1
    zeros = torch.zeros(zero_shape, dtype=x.dtype)
    return torch.cat((x, zeros), dim=dim)