Unverified Commit 0db573c3 authored by Dzianis Dus's avatar Dzianis Dus Committed by GitHub
Browse files

CUDATreeLearner: free GPU memory in destructor if any allocated (#4963)

* CUDATreeLearner: free GPU memory in destruuctor if any allocated

* Minor changes: checking for num_gpu_feature_groups is not needed

* Trigger CI again
parent b99a8f01
......@@ -63,6 +63,43 @@ CUDATreeLearner::CUDATreeLearner(const Config* config)
}
CUDATreeLearner::~CUDATreeLearner() {
#pragma omp parallel for schedule(static, num_gpu_)
for (int device_id = 0; device_id < num_gpu_; ++device_id) {
CUDASUCCESS_OR_FATAL(cudaSetDevice(device_id));
if (device_features_[device_id] != NULL) {
CUDASUCCESS_OR_FATAL(cudaFree(device_features_[device_id]));
}
if (device_gradients_[device_id] != NULL) {
CUDASUCCESS_OR_FATAL(cudaFree(device_gradients_[device_id]));
}
if (device_hessians_[device_id] != NULL) {
CUDASUCCESS_OR_FATAL(cudaFree(device_hessians_[device_id]));
}
if (device_feature_masks_[device_id] != NULL) {
CUDASUCCESS_OR_FATAL(cudaFree(device_feature_masks_[device_id]));
}
if (device_data_indices_[device_id] != NULL) {
CUDASUCCESS_OR_FATAL(cudaFree(device_data_indices_[device_id]));
}
if (sync_counters_[device_id] != NULL) {
CUDASUCCESS_OR_FATAL(cudaFree(sync_counters_[device_id]));
}
if (device_subhistograms_[device_id] != NULL) {
CUDASUCCESS_OR_FATAL(cudaFree(device_subhistograms_[device_id]));
}
if (device_histogram_outputs_[device_id] != NULL) {
CUDASUCCESS_OR_FATAL(cudaFree(device_histogram_outputs_[device_id]));
}
}
}
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment