traced type

EnzymeAD · Jan 8, 2025 · fa2a1df · fa2a1df
1 parent a93e455
commit fa2a1df
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/ext/ReactantCUDAExt.jl b/ext/ReactantCUDAExt.jl
@@ -613,6 +613,12 @@ Reactant.@reactant_overlay @noinline function CUDA.cufunction(
     return Core.Typeof(res)(f, res.entry)
 end
 
+function Reactant.traced_type(
+    ::Type{A}, seen::ST, ::Val{mode}, track_numbers
+) where {T,N,A<:CuTracedArray,ST,mode}
+    return A
+end
+
 function Reactant.traced_type(
     ::Type{A}, seen::ST, ::Val{mode}, track_numbers
 ) where {T,N,A<:CUDA.CuArray{T,N},ST,mode}