diff --git a/gtsam/hybrid/HybridBayesNet.cpp b/gtsam/hybrid/HybridBayesNet.cpp index 841b74f4f..d6fd7e6bd 100644 --- a/gtsam/hybrid/HybridBayesNet.cpp +++ b/gtsam/hybrid/HybridBayesNet.cpp @@ -61,13 +61,15 @@ HybridBayesNet HybridBayesNet::prune(size_t maxNrLeaves) const { } } - // Prune the joint. NOTE: imperative and, again, possibly quite expensive. - DiscreteConditional pruned(joint); - pruned.prune(maxNrLeaves); - - // Create a the result starting with the pruned joint. + // Create the result starting with the pruned joint. HybridBayesNet result; - result.push_back(std::make_shared(pruned)); + result.emplace_shared(joint); + // Prune the joint. NOTE: imperative and, again, possibly quite expensive. + result.back()->asDiscrete()->prune(maxNrLeaves); + + // Get pruned discrete probabilities so + // we can prune HybridGaussianConditionals. + DiscreteConditional pruned = *result.back()->asDiscrete(); /* To prune, we visitWith every leaf in the HybridGaussianConditional. * For each leaf, using the assignment we can check the discrete decision tree