A scheleton of the network is the following:
At the end my network compute the MSE loss of the trajectory and the Kullback Leibler loss (with gaussian prior with mean 0 and std equal to 1) given as:
kld_loss = -0.5*torch.sum(1 + sigma - mu.pow(2) - sigma.exp())
Any idea of the possible causes? Do you need further details?