From 2644c4baa9c07f5d95100f556e3b54ed5a0e4757 Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Fri, 24 Jan 2025 18:04:33 -0500
Subject: [PATCH] prune and then add to BayesNet

---
 gtsam/hybrid/HybridBayesNet.cpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/gtsam/hybrid/HybridBayesNet.cpp b/gtsam/hybrid/HybridBayesNet.cpp
index c95cc686f..77393ab1b 100644
--- a/gtsam/hybrid/HybridBayesNet.cpp
+++ b/gtsam/hybrid/HybridBayesNet.cpp
@@ -58,11 +58,12 @@ HybridBayesNet HybridBayesNet::prune(size_t maxNrLeaves,
     joint = joint * (*conditional);
   }
 
+  // Prune the joint. NOTE: imperative and, again, possibly quite expensive.
+  joint.prune(maxNrLeaves);
+
   // Create the result starting with the pruned joint.
   HybridBayesNet result;
   result.emplace_shared<DiscreteConditional>(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.