diff --git a/mindspore/ccsrc/runtime/device/gpu/gpu_kernel_runtime.cc b/mindspore/ccsrc/runtime/device/gpu/gpu_kernel_runtime.cc index d77c0949da..17abce0c3d 100644 --- a/mindspore/ccsrc/runtime/device/gpu/gpu_kernel_runtime.cc +++ b/mindspore/ccsrc/runtime/device/gpu/gpu_kernel_runtime.cc @@ -53,6 +53,7 @@ using mindspore::device::memswap::MemSwapInfoSet; using mindspore::device::memswap::MemSwapManager; using mindspore::device::memswap::SwapKind; static const size_t PARAMETER_OUTPUT_INDEX = 0; +static thread_local bool cur_thread_device_inited{false}; bool GPUKernelRuntime::SyncStream() { if (!GPUDeviceManager::GetInstance().SyncStream(stream_)) { @@ -71,7 +72,11 @@ bool GPUKernelRuntime::Init() { MS_EXCEPTION_IF_NULL(context_ptr); enable_relation_cache_ = context_ptr->get_param(MS_CTX_ENABLE_GRAPH_KERNEL); - if (device_init_ == true) { + if (device_init_) { + if (!cur_thread_device_inited) { + CHECK_OP_RET_WITH_EXCEPT(CudaDriver::set_current_device(UintToInt(device_id_)), "Failed to set device id"); + cur_thread_device_inited = true; + } GPUMemoryAllocator::GetInstance().CheckMaxDeviceMemory(); return true; }