sync every stream

560d3f1b · Jiezhong Qiu · 610752d2 · 560d3f1b
Commit 560d3f1b authored Dec 30, 2020 by Jiezhong Qiu
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

pytorch/cuda/moe_cuda_kernel.cu pytorch/cuda/moe_cuda_kernel.cu +3 -1

No files found.
--- a/pytorch/cuda/moe_cuda_kernel.cu
+++ b/pytorch/cuda/moe_cuda_kernel.cu
@@ -216,7 +216,9 @@ void moe_cuda_grad_weight(
            grad_weight + gate_host[i] * out_feat * in_feat,
            out_feat));
    }
-    checkCudaErrors(cudaDeviceSynchronize());
+    for (size_t i=0; i<num_expert; ++i) {
+        checkCudaErrors(cudaStreamSynchronize(*(h->streams + i)));
+    }
    delete[] gate_host;
 }