import torch


def reparametisation_trick(mu, log_var, device):
    """

    :param mu: The mean of the latent variable to be formed (nbatch, n_z)
    :param log_var: The log variance of the latent variable to be formed (nbatch, n_z)
    :param device: CPU or GPU
    :return: latent variable (nbatch, n_z)
    """
    noise = torch.normal(mean=0, std=1.0, size=log_var.shape).to(torch.device(device))
    z = mu + torch.mul(torch.exp(log_var / 2.0), noise)

    return z