Fix unified_array

Author: michel2323

ext/OceananigansCUDAExt.jl

@@ -66,6 +66,9 @@ AC.on_architecture(::CUDAGPU, a::SubArray{<:Any, <:Any, <:Array}) = CuArray(a)
 AC.on_architecture(::AC.CPU, a::SubArray{<:Any, <:Any, <:CuArray}) = Array(a)
 AC.on_architecture(::CUDAGPU, a::StepRangeLen) = a
 
+# cu alters the type of `a`, so we convert it back to the correct type
+unified_array(::CUDAGPU, a::AbstractArray) = map(eltype(a), cu(a; unified = true))
+
 ## GPU to GPU copy of contiguous data
 @inline function AC.device_copy_to!(dst::CuArray, src::CuArray; async::Bool = false)
     n = length(src)

src/Architectures.jl

@@ -104,9 +104,6 @@ cpu_architecture(::ReactantState) = CPU()
 unified_array(::CPU, a) = a
 unified_array(::GPU, a) = a
 
-# cu alters the type of `a`, so we convert it back to the correct type
-unified_array(::GPU, a::AbstractArray) = map(eltype(a), cu(a; unified = true))
-
 @inline device_copy_to!(dst::Array, src::Array; kw...) = Base.copyto!(dst, src)
 
 @inline unsafe_free!(a)          = nothing