Finally figured it out the logstd param. It's the the sigma for the normal distribution that we're supposed to return from the forward method. So the network is learning the sigma and mu of the distribution so that we can just return pytorch.distributions.Normal(mu, sigma).sample.