Hotfix for new OneElement on GPU

FluxML · Jan 8, 2024 · ea3246c · ea3246c
1 parent 1af3f4d
commit ea3246c
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/ext/FluxCUDAExt/functor.jl b/ext/FluxCUDAExt/functor.jl
@@ -29,6 +29,10 @@ adapt_storage(to::FluxCUDAAdaptor, x::AbstractRNG) =
 
 # TODO: figure out the correct design for OneElement
 adapt_storage(to::FluxCUDAAdaptor, x::Zygote.OneElement) = CUDA.cu(collect(x))
+# Patch for GPU support until we can make OneElement smarter
+if isdefined(Zygote.ChainRules, :OneElement)
+    adapt_storage(to::FluxCUDAAdaptor, x::Zygote.ChainRules.OneElement) = CUDA.cu(collect(x))
+end
 
 adapt_storage(to::FluxCPUAdaptor, x::T) where T <: CUDA.CUSPARSE.CUDA.CUSPARSE.AbstractCuSparseMatrix = adapt(Array, x)
 adapt_storage(to::FluxCPUAdaptor, x::CUDA.RNG) = Random.default_rng()