import torch


def pad_tensor(vec, pad, dim):
    """
        Pads a tensor with zeros according to arguments given

        Args:
            vec (Tensor): Tensor to pad
            pad (int): The total tensor size with pad
            dim (int): Dimension to pad

        Returns:
            padded_tensor (Tensor): A new tensor padded to 'pad' in dimension 'dim'

    """
    pad_size = list(vec.shape)
    pad_size[dim] = pad - vec.size(dim)
    padded_tensor = torch.cat([vec, torch.zeros(*pad_size)], dim=dim)
    return padded_tensor