Sync cuda stream before updating stats

Signed-off-by: Michael Tuttle <[email protected]>
quic · Oct 17, 2023 · 9890b66 · 9890b66
1 parent 2be0b4e
commit 9890b66
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/TrainingExtensions/onnx/src/QcQuantizeOp.cpp b/TrainingExtensions/onnx/src/QcQuantizeOp.cpp
@@ -88,6 +88,12 @@ void QcQuantizeKernel::Compute(OrtKernelContext* context)
     {
         allocator = &cudaAllocator;
         stream = api_.KernelContext_GetGPUComputeStream(context);
+        if ((opMode == DlQuantization::TensorQuantizerOpMode::updateStats) ||
+            (opMode == DlQuantization::TensorQuantizerOpMode::oneShotQuantizeDequantize))
+        {
+            // updateStats doesn't use cuda stream, must synchronize first to ensure input buffer is populated
+            cudaStreamSynchronize(reinterpret_cast<cudaStream_t>(stream));
+        }
     }
 #endif