diff --git a/gtsam/hybrid/HybridFactorGraph.cpp b/gtsam/hybrid/HybridFactorGraph.cpp
index a395e17018..78c051b17f 100644
--- a/gtsam/hybrid/HybridFactorGraph.cpp
+++ b/gtsam/hybrid/HybridFactorGraph.cpp
@@ -18,6 +18,7 @@
  */
 
 #include <gtsam/hybrid/HybridFactorGraph.h>
+#include <gtsam/nonlinear/NonlinearFactor.h>
 
 namespace gtsam {
 
@@ -58,6 +59,8 @@ const KeySet HybridFactorGraph::continuousKeySet() const {
       }
     } else if (auto p = std::dynamic_pointer_cast<GaussianFactor>(factor)) {
       keys.insert(p->keys().begin(), p->keys().end());
+    } else if (auto p = std::dynamic_pointer_cast<NonlinearFactor>(factor)) {
+      keys.insert(p->keys().begin(), p->keys().end());
     }
   }
   return keys;
diff --git a/gtsam/hybrid/HybridGaussianConditional.cpp b/gtsam/hybrid/HybridGaussianConditional.cpp
index 6bd859c1da..4e896206f7 100644
--- a/gtsam/hybrid/HybridGaussianConditional.cpp
+++ b/gtsam/hybrid/HybridGaussianConditional.cpp
@@ -69,12 +69,14 @@ HybridGaussianConditional::HybridGaussianConditional(
       conditionals_(conditionals),
       negLogConstant_(helper.minNegLogConstant) {}
 
+/* *******************************************************************************/
 HybridGaussianConditional::HybridGaussianConditional(
     const DiscreteKeys &discreteParents,
     const HybridGaussianConditional::Conditionals &conditionals)
     : HybridGaussianConditional(discreteParents, conditionals,
                                 ConstructorHelper(conditionals)) {}
 
+/* *******************************************************************************/
 HybridGaussianConditional::HybridGaussianConditional(
     const DiscreteKey &discreteParent,
     const std::vector<GaussianConditional::shared_ptr> &conditionals)
@@ -242,13 +244,6 @@ std::set<DiscreteKey> DiscreteKeysAsSet(const DiscreteKeys &discreteKeys) {
 }
 
 /* ************************************************************************* */
-/**
- * @brief Helper function to get the pruner functional.
- *
- * @param discreteProbs The probabilities of only discrete keys.
- * @return std::function<GaussianConditional::shared_ptr(
- * const Assignment<Key> &, const GaussianConditional::shared_ptr &)>
- */
 std::function<GaussianConditional::shared_ptr(
     const Assignment<Key> &, const GaussianConditional::shared_ptr &)>
 HybridGaussianConditional::prunerFunc(const DecisionTreeFactor &discreteProbs) {
diff --git a/gtsam/hybrid/HybridGaussianConditional.h b/gtsam/hybrid/HybridGaussianConditional.h
index e0bf66e529..02f6bba1ea 100644
--- a/gtsam/hybrid/HybridGaussianConditional.h
+++ b/gtsam/hybrid/HybridGaussianConditional.h
@@ -194,7 +194,13 @@ class GTSAM_EXPORT HybridGaussianConditional
   /// Convert to a DecisionTree of Gaussian factor graphs.
   GaussianFactorGraphTree asGaussianFactorGraphTree() const;
 
-  //// Get the pruner functor from pruned discrete probabilities.
+  /**
+   * @brief Get the pruner function from discrete probabilities.
+   *
+   * @param discreteProbs The probabilities of only discrete keys.
+   * @return std::function<GaussianConditional::shared_ptr(
+   * const Assignment<Key> &, const GaussianConditional::shared_ptr &)>
+   */
   std::function<GaussianConditional::shared_ptr(
       const Assignment<Key> &, const GaussianConditional::shared_ptr &)>
   prunerFunc(const DecisionTreeFactor &prunedProbabilities);
diff --git a/gtsam/hybrid/HybridGaussianFactor.cpp b/gtsam/hybrid/HybridGaussianFactor.cpp
index 6f33c53284..7f8e808bf4 100644
--- a/gtsam/hybrid/HybridGaussianFactor.cpp
+++ b/gtsam/hybrid/HybridGaussianFactor.cpp
@@ -126,16 +126,19 @@ HybridGaussianFactor::HybridGaussianFactor(const ConstructorHelper &helper)
       factors_(helper.factorsTree.empty() ? augment(helper.pairs)
                                           : helper.factorsTree) {}
 
+/* *******************************************************************************/
 HybridGaussianFactor::HybridGaussianFactor(
     const DiscreteKey &discreteKey,
     const std::vector<GaussianFactor::shared_ptr> &factors)
     : HybridGaussianFactor(ConstructorHelper(discreteKey, factors)) {}
 
+/* *******************************************************************************/
 HybridGaussianFactor::HybridGaussianFactor(
     const DiscreteKey &discreteKey,
     const std::vector<GaussianFactorValuePair> &factorPairs)
     : HybridGaussianFactor(ConstructorHelper(discreteKey, factorPairs)) {}
 
+/* *******************************************************************************/
 HybridGaussianFactor::HybridGaussianFactor(const DiscreteKeys &discreteKeys,
                                            const FactorValuePairs &factors)
     : HybridGaussianFactor(ConstructorHelper(discreteKeys, factors)) {}
@@ -223,6 +226,7 @@ double HybridGaussianFactor::potentiallyPrunedComponentError(
     return std::numeric_limits<double>::max();
   }
 }
+
 /* *******************************************************************************/
 AlgebraicDecisionTree<Key> HybridGaussianFactor::errorTree(
     const VectorValues &continuousValues) const {