import torch


def sample_z(z_mean, z_std, how_many, device=None):
    """
    Returns a sample from z of size (z_dim, how_many)
    """
    z_dim = z_std.shape[0]
    if device:
        eps = torch.randn([z_dim, how_many]).to(device)
    else:
        eps = torch.randn([z_dim, how_many])
    z_samples = z_mean + z_std * eps
    return z_samples