未验证 提交 89e54d69 编写于 作者: L Leo Chen 提交者: GitHub

skip first time reset (#55498)

上级 8e9e0659
......@@ -30,8 +30,21 @@ class CUDADeviceGuard {
CUDADeviceGuard() {}
~CUDADeviceGuard() {
static thread_local bool is_first_time_ = true;
if (prev_id_ != -1) {
// 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;
}
}
}
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册