Merge pull request #19 from darsnack/kd/fix-@atomic

Use `CUDA.@atomic`
FluxML · Jul 2, 2021 · af81ef6 · af81ef6
2 parents 7577b98 + 2ce8a27
commit af81ef6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ext/CUDAExt/src/utils.jl b/ext/CUDAExt/src/utils.jl
@@ -19,7 +19,7 @@ function divide_kernel!(xs, counts, max_idx, max_dims_idx, dims_size)
     @inbounds if index <= max_idx
         j, k = divrem(index-1, max_dims_idx)
         dims_i = Tuple(CartesianIndices(dims_size)[k+1])
-        @atomic xs[dims_i..., j+1] = xs[dims_i..., j+1] / counts[j+1]
+        CUDA.@atomic xs[dims_i..., j+1] = xs[dims_i..., j+1] / counts[j+1]
     end
     return nothing
 end