From a725272e19748a5f3ddd7eeeff7c766660e38a06 Mon Sep 17 00:00:00 2001
From: Rory Mitchell <r.a.mitchell.nz@gmail.com>
Date: Tue, 24 Jul 2018 15:03:36 +1200
Subject: [PATCH] Correct mistake from dmatrix refactor (#3408)

---
 src/tree/updater_histmaker.cc | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/tree/updater_histmaker.cc b/src/tree/updater_histmaker.cc
index 3ab48e2a8..3cc0a760b 100644
--- a/src/tree/updater_histmaker.cc
+++ b/src/tree/updater_histmaker.cc
@@ -434,9 +434,10 @@ class CQHistMaker: public HistMaker<TStats> {
         this->CorrectNonDefaultPositionByBatch(batch, fsplit_set_, tree);
 
         // start enumeration
-        const auto nsize = static_cast<bst_omp_uint>(batch.Size());
+        const auto nsize = static_cast<bst_omp_uint>(work_set_.size());
         #pragma omp parallel for schedule(dynamic, 1)
-        for (bst_omp_uint fid = 0; fid < nsize; ++fid) {
+        for (bst_omp_uint i = 0; i < nsize; ++i) {
+          int fid = work_set_[i];
           int offset = feat2workindex_[fid];
           if (offset >= 0) {
             this->UpdateSketchCol(gpair, batch[fid], tree,