Put stream checking out of callback

tensorflow · Aug 24, 2023 · 332e1fe · 332e1fe
1 parent ccae79b
commit 332e1fe
Showing 1 changed file with 11 additions and 4 deletions.
diff --git a/tensorflow/core/common_runtime/gpu/gpu_util.cc b/tensorflow/core/common_runtime/gpu/gpu_util.cc
@@ -397,7 +397,12 @@ void GPUUtil::CopyCPUTensorToGPU(const Tensor* cpu_tensor,
     }
   }
 
-  if (merge_host_to_device_stream) done(OkStatus());
+  if (merge_host_to_device_stream) {
+    if (!recv_host_to_device_stream->ok()) {
+      LOG(FATAL) << "CPU->GPU Memcpy failed";
+    }
+    done(OkStatus());
+  }
   dev_info->event_mgr->ThenExecute(
       recv_host_to_device_stream,
       [recv_host_to_device_stream, done, input_ref, do_staging, staging_buffer,
@@ -407,10 +412,12 @@ void GPUUtil::CopyCPUTensorToGPU(const Tensor* cpu_tensor,
         } else {
           input_ref.Unref();
         }
-        if (!recv_host_to_device_stream->ok()) {
-          LOG(FATAL) << "CPU->GPU Memcpy failed";
+        if (!merge_host_to_device_stream) {
+          if (!recv_host_to_device_stream->ok()) {
+            LOG(FATAL) << "CPU->GPU Memcpy failed";
+          }
+          done(OkStatus());
         }
-        if (!merge_host_to_device_stream) done(OkStatus());
       });
 }