import torch


def pad(data_tensor, pad_size, dim, pad_at_end=True):
    """
    Pads a Tensor with zeros along a dimension.

    :param data_tensor: Tensor to pad.
    :param pad_size: How many zeros to append.
    :param dim: The dimension to pad.
    :param pad_at_end: Whether to pad at the end.
    :return: Padded Tensor.
    """

    if dim < 0:
        dim = len(data_tensor.shape) + dim

    shape = list(data_tensor.shape)
    shape[dim] = pad_size

    zeros = torch.zeros(shape, device=data_tensor.device)

    if pad_at_end:
        return torch.cat([data_tensor, zeros], dim=dim)
    else:
        return torch.cat([zeros, data_tensor], dim=dim)