Turn off preconditioner when computing the KL divergence
@@ -73,6 +73,7 @@ class VariationalStrategy(Module):
variational_dist_u = self.variational_distribution.variational_distribution
prior_dist = self.prior_distribution
+ with settings.max_preconditioner_size(0):
kl_divergence = torch.distributions.kl.kl_divergence(variational_dist_u, prior_dist)
return kl_divergence
