adding pre pack sync

microsoft · Oct 4, 2023 · 44b7d26 · 44b7d26
1 parent 4630b6d
commit 44b7d26
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/onnxruntime/core/providers/cuda/nn/conv.cc b/onnxruntime/core/providers/cuda/nn/conv.cc
@@ -118,6 +118,7 @@ Status Conv<T, NHWC>::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr
       if (!status.IsOK()) {
         return status;
       }
+      CUDA_CALL_THROW(cudaStreamSynchronize(DefaultCudaStream()));
       is_packed = true;
     }
   }

diff --git a/onnxruntime/core/providers/cuda/nn/conv_transpose.cc b/onnxruntime/core/providers/cuda/nn/conv_transpose.cc
@@ -77,6 +77,7 @@ Status ConvTranspose<T, NHWC>::PrePack(const Tensor& tensor, int input_idx, Allo
       if (!status.IsOK()) {
         return status;
       }
+      CUDA_CALL_THROW(cudaStreamSynchronize(DefaultCudaStream()));
       is_packed = true;
     }
   }