Modify caching allocator/vector and fix issues relating to inability to train large datasets (#4615)

This commit is contained in:
sriramch
2019-07-08 23:33:27 -07:00
committed by Rory Mitchell
parent cd1526d3b1
commit 7a388cbf8b
5 changed files with 22 additions and 18 deletions

View File

@@ -250,6 +250,10 @@ class GPUPredictor : public xgboost::Predictor {
struct DeviceShard {
DeviceShard() : device_{-1} {}
~DeviceShard() {
dh::safe_cuda(cudaSetDevice(device_));
}
void Init(int device) {
this->device_ = device;
max_shared_memory_bytes_ = dh::MaxSharedMemory(this->device_);