diff --git a/gtsam/discrete/DecisionTreeFactor.cpp b/gtsam/discrete/DecisionTreeFactor.cpp
index 4e16fc689..7f604086c 100644
--- a/gtsam/discrete/DecisionTreeFactor.cpp
+++ b/gtsam/discrete/DecisionTreeFactor.cpp
@@ -156,9 +156,9 @@ namespace gtsam {
   std::vector<std::pair<DiscreteValues, double>> DecisionTreeFactor::enumerate()
       const {
     // Get all possible assignments
-    std::vector<std::pair<Key, size_t>> pairs = discreteKeys();
+    DiscreteKeys pairs = discreteKeys();
     // Reverse to make cartesian product output a more natural ordering.
-    std::vector<std::pair<Key, size_t>> rpairs(pairs.rbegin(), pairs.rend());
+    DiscreteKeys rpairs(pairs.rbegin(), pairs.rend());
     const auto assignments = DiscreteValues::CartesianProduct(rpairs);
 
     // Construct unordered_map with values
diff --git a/gtsam/hybrid/HybridDiscreteFactor.cpp b/gtsam/hybrid/HybridDiscreteFactor.cpp
index dd8b867be..afdb6472a 100644
--- a/gtsam/hybrid/HybridDiscreteFactor.cpp
+++ b/gtsam/hybrid/HybridDiscreteFactor.cpp
@@ -26,7 +26,6 @@
 namespace gtsam {
 
 /* ************************************************************************ */
-// TODO(fan): THIS IS VERY VERY DIRTY! We need to get DiscreteFactor right!
 HybridDiscreteFactor::HybridDiscreteFactor(DiscreteFactor::shared_ptr other)
     : Base(boost::dynamic_pointer_cast<DecisionTreeFactor>(other)
                ->discreteKeys()),