import torch


def get_noise(n_samples, z_dim, device='cpu'):
    """Create noise vectors.

    Given the dimensions (n_samples, z_dim), creates a tensor of that shape
    filled with random numbers from the normal distribution.

    Args:
        n_samples (int): the number of samples to generate, a scalar
        z_dim (int): the dimension of the noise vector, a scalar
        device (device): the device type

    Returns:
        Tensor: random vector
    """
    # NOTE: To use this on GPU with device='cuda', make sure to pass the device
    # argument to the function you use to generate the noise.
    return torch.randn(n_samples, z_dim).to(device)