diff --git a/gtsam/hybrid/HybridJunctionTree.cpp b/gtsam/hybrid/HybridJunctionTree.cpp
index 573df7eca..a463c625f 100644
--- a/gtsam/hybrid/HybridJunctionTree.cpp
+++ b/gtsam/hybrid/HybridJunctionTree.cpp
@@ -62,8 +62,12 @@ struct HybridConstructorTraversalData {
 
     // Add all the discrete keys in the hybrid factors to the current data
     for (const auto& f : node->factors) {
-      if (auto p = boost::dynamic_pointer_cast<HybridFactor>(f)) {
-        for (auto& k : p->discreteKeys()) {
+      if (auto hf = boost::dynamic_pointer_cast<HybridFactor>(f)) {
+        for (auto& k : hf->discreteKeys()) {
+          data.discreteKeys.insert(k.first);
+        }
+      } else if (auto hf = boost::dynamic_pointer_cast<DecisionTreeFactor>(f)) {
+        for (auto& k : hf->discreteKeys()) {
           data.discreteKeys.insert(k.first);
         }
       }