From 89e54d69926b60fd59ae52bcad02abe111c104d1 Mon Sep 17 00:00:00 2001 From: Leo Chen Date: Wed, 19 Jul 2023 17:15:52 +0800 Subject: [PATCH] skip first time reset (#55498) --- paddle/fluid/platform/cuda_device_guard.h | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/paddle/fluid/platform/cuda_device_guard.h b/paddle/fluid/platform/cuda_device_guard.h index 08beed532a7..53618286040 100644 --- a/paddle/fluid/platform/cuda_device_guard.h +++ b/paddle/fluid/platform/cuda_device_guard.h @@ -30,8 +30,21 @@ class CUDADeviceGuard { CUDADeviceGuard() {} ~CUDADeviceGuard() { + static thread_local bool is_first_time_ = true; if (prev_id_ != -1) { - platform::SetDeviceId(prev_id_); + // Do not set device back for the first time, since + // `cudaGetDevice` returns 0 when `cudaSetDevice` is + // not called. + // In that case, if CUDADeviceGuard(7) is called, + // prev_id will be 0 and we don`t need to set it back to 0. + // If cudaSetDevice(0) is called, it may use hundreds MB of + // the gpu memory. + VLOG(10) << __func__ << " prev_id: " << prev_id_ << ", is_first_time_" + << is_first_time_; + if (!(is_first_time_ && prev_id_ == 0)) { + platform::SetDeviceId(prev_id_); + is_first_time_ = false; + } } } -- GitLab