From fa542a20384435c88eaf3cf2b460a2c645908812 Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Fri, 18 Mar 2022 08:19:01 -0400
Subject: [PATCH] address review comments

---
 gtsam/discrete/DecisionTree-inl.h         | 8 +++++---
 gtsam/discrete/tests/testDecisionTree.cpp | 6 ++++--
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/gtsam/discrete/DecisionTree-inl.h b/gtsam/discrete/DecisionTree-inl.h
index c616c0269..627c1a5aa 100644
--- a/gtsam/discrete/DecisionTree-inl.h
+++ b/gtsam/discrete/DecisionTree-inl.h
@@ -668,9 +668,11 @@ namespace gtsam {
       for (size_t i = 0; i < choice->nrChoices(); i++) {
         choices[choice->label()] = i;  // Set assignment for label to i
 
-        VisitWith<L, Y> visit(f);
-        visit.choices = choices;
-        visit(choice->branches()[i]);  // recurse!
+        (*this)(choice->branches()[i]);  // recurse!
+
+        // Remove the choice so we are backtracking
+        auto choice_it = choices.find(choice->label());
+        choices.erase(choice_it);
       }
     }
   };
diff --git a/gtsam/discrete/tests/testDecisionTree.cpp b/gtsam/discrete/tests/testDecisionTree.cpp
index ca5daf201..91deed625 100644
--- a/gtsam/discrete/tests/testDecisionTree.cpp
+++ b/gtsam/discrete/tests/testDecisionTree.cpp
@@ -20,11 +20,13 @@
 // #define DT_DEBUG_MEMORY
 // #define DT_NO_PRUNING
 #define DISABLE_DOT
-#include <CppUnitLite/TestHarness.h>
-#include <gtsam/base/Testable.h>
 #include <gtsam/discrete/DecisionTree-inl.h>
+
+#include <gtsam/base/Testable.h>
 #include <gtsam/discrete/Signature.h>
 
+#include <CppUnitLite/TestHarness.h>
+
 #include <boost/assign/std/vector.hpp>
 using namespace boost::assign;