1 files changed, 4 insertions, 1 deletions
diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.cc b/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.cc
index 881e4339a7..952b34b353 100644
--- a/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.cc
+++ b/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.cc
@@ -23,12 +23,15 @@ using decision_trees::TreeNode;
 
 int32 DecisionTreeResource::TraverseTree(
     const std::unique_ptr<TensorDataSet>& input_data, int example,
-    int32* leaf_depth) const {
+    int32* leaf_depth, TreePath* path) const {
   const DecisionTree& tree = decision_tree_->decision_tree();
   int32 current_id = 0;
   int32 depth = 0;
   while (true) {
     const TreeNode& current = tree.nodes(current_id);
+    if (path != nullptr) {
+      *path->add_nodes_visited() = current;
+    }
     if (current.has_leaf()) {
       if (leaf_depth != nullptr) {
         *leaf_depth = depth;