diff --git a/src/io/page_fmatrix-inl.hpp b/src/io/page_fmatrix-inl.hpp
index 971abbb0e..91d24cf2d 100644
--- a/src/io/page_fmatrix-inl.hpp
+++ b/src/io/page_fmatrix-inl.hpp
@@ -114,11 +114,11 @@ class CSCMatrixManager {
   }
   inline void SetColSet(const std::vector<bst_uint> &cset, bool setall) {
     if (!setall) {
-      col_todo_.resize(cset.size());
+      col_todo_.resize(0);
       for (size_t i = 0; i < cset.size(); ++i) {
-        col_todo_[i] = cset[i];
-        utils::Assert(col_todo_[i] < static_cast<bst_uint>(col_ptr_.size() - 1),
-                      "CSCMatrixManager: column index exceed bound");
+        if (col_todo_[i] < static_cast<bst_uint>(col_ptr_.size() - 1)) {
+          col_todo_.push_back(cset[i]);
+        }
       }
       std::sort(col_todo_.begin(), col_todo_.end());
     } else {
diff --git a/src/io/simple_fmatrix-inl.hpp b/src/io/simple_fmatrix-inl.hpp
index 88bc69019..08e25e28b 100644
--- a/src/io/simple_fmatrix-inl.hpp
+++ b/src/io/simple_fmatrix-inl.hpp
@@ -76,7 +76,11 @@ class FMatrixS : public IFMatrix{
    * \brief colmun based iterator
    */
   virtual utils::IIterator<ColBatch> *ColIterator(const std::vector<bst_uint> &fset) {
-    col_iter_.col_index_ = fset;
+    size_t ncol = this->NumCol();
+    col_iter_.col_index_.resize(0);
+    for (size_t i = 0; i < fset.size(); ++i) {
+      if (fset[i] < ncol) col_iter_.col_index_.push_back(fset[i]); 
+    }
     col_iter_.SetBatch(col_ptr_, col_data_);
     return &col_iter_;
   }