Fixed asynchronous streams in column reduce kernel call

2126b59a · TiagoMAntunes · c96f8863 · 2126b59a
Commit 2126b59a authored Jun 09, 2021 by TiagoMAntunes
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

cuda/parallel_linear.cuh cuda/parallel_linear.cuh +1 -1

No files found.
--- a/cuda/parallel_linear.cuh
+++ b/cuda/parallel_linear.cuh
@@ -146,7 +146,7 @@ void fmoe_cuda_linear_backward_impl(
        
        if (has_bias) {
            column_reduce
-            <<<grid_threads, block_threads, sizeof(scalar_t)*1024, smgr->stream(0)>>>
+            <<<grid_threads, block_threads, sizeof(scalar_t)*1024, smgr->stream(i)>>>
            (
                grad_output_buf + ptr * out_feat,
                grad_bias + i * out_feat,