allow more flexible proposals in atomic snpe-c

sbi-dev · Aug 24, 2022 · ce42bad · ce42bad
1 parent 893ac43
commit ce42bad
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/sbi/inference/snpe/snpe_c.py b/sbi/inference/snpe/snpe_c.py
@@ -161,7 +161,8 @@ def train(
             # last proposal.
             proposal = self._proposal_roundwise[-1]
             self.use_non_atomic_loss = (
-                isinstance(proposal.posterior_estimator._distribution, mdn)
+                isinstance(proposal, DirectPosterior)
+                and isinstance(proposal.posterior_estimator._distribution, mdn)
                 and isinstance(self._neural_net._distribution, mdn)
                 and check_dist_class(
                     self._prior, class_to_check=(Uniform, MultivariateNormal)