support SubArray on fill

CliMA · Jul 20, 2023 · 0dd6265 · 0dd6265
1 parent 6af8766
commit 0dd6265
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/DataLayouts/cuda.jl b/src/DataLayouts/cuda.jl
@@ -97,7 +97,7 @@ end
 function Base.fill!(
     dest::IJFH{S, Nij, A},
     val,
-) where {S, Nij, A <: CUDA.CuArray}
+) where {S, Nij, A <: Union{CUDA.CuArray, SubArray{<:Any, <:Any, CUDA.CuArray}}}
     _, _, _, _, Nh = size(dest)
     if Nh > 0
         CUDA.@cuda threads = (Nij, Nij) blocks = (Nh, 1) knl_fill!(dest, val)