def forward(self, hidden_states, attention_mask):
    latent_states = hidden_states[:, 0, :]
    mus = self.mu_embedding(latent_states)
    sigmas = self.sigma_embedding(latent_states)
    sigmas = torch.exp((sigmas / 2))
    random_normal = self.gaussian_generator.sample([sigmas.size(0)]).to(sigmas.device)
    zs = (mus + (sigmas * random_normal))
    return (mus, sigmas, zs)
def forward(self, hidden_states, attention_mask):
    latent_states = hidden_states[:, 0, :]
    mus = self.mu_embedding(latent_states)
    sigmas = self.sigma_embedding(latent_states)
    sigmas = torch.exp((sigmas / 2))
    random_normal = self.gaussian_generator.sample([sigmas.size(0)]).to(sigmas.device)
    zs = (mus + (sigmas * random_normal))
    return (mus, sigmas, zs)