diff --git a/gtsam/hybrid/GaussianMixture.cpp b/gtsam/hybrid/GaussianMixture.cpp
index 79c385262..753e35bf0 100644
--- a/gtsam/hybrid/GaussianMixture.cpp
+++ b/gtsam/hybrid/GaussianMixture.cpp
@@ -234,7 +234,7 @@ std::set<DiscreteKey> DiscreteKeysAsSet(const DiscreteKeys &discreteKeys) {
  */
 std::function<GaussianConditional::shared_ptr(
     const Assignment<Key> &, const GaussianConditional::shared_ptr &)>
-GaussianMixture::prunerFunc(const TableFactor &discreteProbs) {
+GaussianMixture::prunerFunc(const DecisionTreeFactor &discreteProbs) {
   // Get the discrete keys as sets for the decision tree
   // and the gaussian mixture.
   auto discreteProbsKeySet = DiscreteKeysAsSet(discreteProbs.discreteKeys());
@@ -285,7 +285,7 @@ GaussianMixture::prunerFunc(const TableFactor &discreteProbs) {
 }
 
 /* *******************************************************************************/
-void GaussianMixture::prune(const TableFactor &discreteProbs) {
+void GaussianMixture::prune(const DecisionTreeFactor &discreteProbs) {
   // Functional which loops over all assignments and create a set of
   // GaussianConditionals
   auto pruner = prunerFunc(discreteProbs);
diff --git a/gtsam/hybrid/GaussianMixture.h b/gtsam/hybrid/GaussianMixture.h
index c193d065c..0b68fcfd0 100644
--- a/gtsam/hybrid/GaussianMixture.h
+++ b/gtsam/hybrid/GaussianMixture.h
@@ -23,7 +23,6 @@
 #include <gtsam/discrete/DecisionTree.h>
 #include <gtsam/discrete/DecisionTreeFactor.h>
 #include <gtsam/discrete/DiscreteKey.h>
-#include <gtsam/discrete/TableFactor.h>
 #include <gtsam/hybrid/GaussianMixtureFactor.h>
 #include <gtsam/hybrid/HybridFactor.h>
 #include <gtsam/inference/Conditional.h>
@@ -81,7 +80,7 @@ class GTSAM_EXPORT GaussianMixture
    */
   std::function<GaussianConditional::shared_ptr(
       const Assignment<Key> &, const GaussianConditional::shared_ptr &)>
-  prunerFunc(const TableFactor &discreteProbs);
+  prunerFunc(const DecisionTreeFactor &discreteProbs);
 
  public:
   /// @name Constructors
@@ -239,7 +238,7 @@ class GTSAM_EXPORT GaussianMixture
    *
    * @param discreteProbs A pruned set of probabilities for the discrete keys.
    */
-  void prune(const TableFactor &discreteProbs);
+  void prune(const DecisionTreeFactor &discreteProbs);
 
   /**
    * @brief Merge the Gaussian Factor Graphs in `this` and `sum` while
diff --git a/gtsam/hybrid/HybridBayesNet.cpp b/gtsam/hybrid/HybridBayesNet.cpp
index 302ccfd6c..266e02b0d 100644
--- a/gtsam/hybrid/HybridBayesNet.cpp
+++ b/gtsam/hybrid/HybridBayesNet.cpp
@@ -64,7 +64,7 @@ DecisionTreeFactor::shared_ptr HybridBayesNet::discreteConditionals() const {
  * @return std::function<double(const Assignment<Key> &, double)>
  */
 std::function<double(const Assignment<Key> &, double)> prunerFunc(
-    const TableFactor &prunedDiscreteProbs,
+    const DecisionTreeFactor &prunedDiscreteProbs,
     const HybridConditional &conditional) {
   // Get the discrete keys as sets for the decision tree
   // and the Gaussian mixture.
diff --git a/gtsam/hybrid/HybridBayesNet.h b/gtsam/hybrid/HybridBayesNet.h
index a71768bac..825f3acfd 100644
--- a/gtsam/hybrid/HybridBayesNet.h
+++ b/gtsam/hybrid/HybridBayesNet.h
@@ -17,6 +17,7 @@
 
 #pragma once
 
+#include <gtsam/discrete/DecisionTreeFactor.h>
 #include <gtsam/global_includes.h>
 #include <gtsam/hybrid/HybridConditional.h>
 #include <gtsam/hybrid/HybridValues.h>
@@ -225,7 +226,8 @@ class GTSAM_EXPORT HybridBayesNet : public BayesNet<HybridConditional> {
    *
    * @param prunedDiscreteProbs
    */
-  void updateDiscreteConditionals(const DecisionTreeFactor &prunedDiscreteProbs);
+  void updateDiscreteConditionals(
+      const DecisionTreeFactor &prunedDiscreteProbs);
 
 #ifdef GTSAM_ENABLE_BOOST_SERIALIZATION
   /** Serialization function */