Fix bug of wrong cuda device id

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<bool>(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;
   }