Use CUDA.@atomic

FluxML · Jul 2, 2021 · 2ce8a27 · 2ce8a27
1 parent 7577b98
commit 2ce8a27
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