       std::copy(src_nodes.begin(), src_nodes.end(),
                 h_nodes.begin() + h_tree_segments[tree_idx - tree_begin]);
     }
+    dh::ExecuteIndexShards(&shards_, [&](int idx, DeviceShard &shard) {
+      shard.InitModel(model, h_tree_segments, h_nodes, tree_begin, tree_end);
+    });
+  }
+
+  void DevicePredictInternal(DMatrix* dmat,
+                             HostDeviceVector<bst_float>* out_preds,
+                             const gbm::GBTreeModel& model, size_t tree_begin,
+                             size_t tree_end) {
+    if (tree_end - tree_begin == 0) { return; }
+    monitor_.StartCuda("DevicePredictInternal");
+
+    InitModel(model, tree_begin, tree_end);
 
     size_t batch_offset = 0;
     for (auto &batch : dmat->GetRowBatches()) {

    thrust::device_vector<size_t> tree_segments_;
    thrust::device_vector<int> tree_group_;
    size_t max_shared_memory_bytes_;
  };

  void DevicePredictInternal(DMatrix* dmat,
                             HostDeviceVector<bst_float>* out_preds,
                             const gbm::GBTreeModel& model, size_t tree_begin,
                             size_t tree_end) {
    if (tree_end - tree_begin == 0) { return; }
    monitor_.StartCuda("DevicePredictInternal");

    CHECK_EQ(model.param.size_leaf_vector, 0);
    // Copy decision trees to device
    thrust::host_vector<size_t> h_tree_segments;

      DeviceOffsets(batch.offset, batch.data.Size(), &device_offsets);
      batch.data.Reshard(GPUDistribution::Explicit(devices_, device_offsets));

      dh::ExecuteIndexShards(&shards_, [&](int idx, DeviceShard& shard) {
        shard.PredictInternal(batch, dmat->Info(), out_preds);
      });
      batch_offset += batch.Size() * model.param.num_output_group;
    }
