Fix __shfl_down cuda error

HawkAaron · Nov 1, 2019 · e950645 · e950645
1 parent 869e473
commit e950645
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -45,7 +45,7 @@ class RNNTLoss(Loss):
 From the repo test with:
 
 ```bash
-python test/test.py 10 300 100 50 --mx
+python test/test.py --mx
 ```
 
 ## Reference

diff --git a/rnnt_include/detail/reduce.h b/rnnt_include/detail/reduce.h
@@ -30,7 +30,7 @@ struct CTAReduce {
 
         T shuff;
         for (int offset = warp_size / 2; offset > 0; offset /= 2) {
-            shuff = __shfl_down(x, offset);
+            shuff = __shfl_down_sync(0xFFFFFFFF, x, offset);
             if (tid + offset < count && tid < offset)
                 x = g(x, shuff);
         }