How to properly get a custom layer on gpu in Julia Flux? - julia

Question

As mentioned, I want to get GPU support for a custom layer. Here is the build-net part of my code:
# MyDense
struct MyDense
d1
d2
end
MyDense(p::Pair) = MyDense(Dense(p), Dense(p))
(m::MyDense)(x) = m.d1(x) + m.d2(x)
# multiple-dispatch for gpu function on the custom layer
Flux.gpu(layer::MyDense) = MyDense(Flux.gpu(layer.d1), Flux.gpu(layer.d2))
# Define our model, a multi-layer perceptron with one hidden layer of size 3:
model = Chain(
Dense(2 => 3, tanh), # activation function inside layer
BatchNorm(3),
MyDense(3 => 2),
softmax) |> gpu # move model to GPU, if available
Reminded by #eliassno's answer in this question, I wrote a new dispatch of Flux.gpu function for my custom layer. However, it still didn't work and gave me the error:
ERROR: ArgumentError: cannot take the CPU address of a CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}
Stacktrace:
[1] unsafe_convert(#unused#::Type{Ptr{Float32}}, x::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
# CUDA C:\Users\Herr LU\.julia\packages\CUDA\DfvRa\src\array.jl:319
[2] gemm!(transA::Char, transB::Char, alpha::Float32, A::Matrix{Float32}, B::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, beta::Float32, C::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
# LinearAlgebra.BLAS C:\Users\Herr LU\AppData\Local\Programs\Julia-1.8.2\share\julia\stdlib\v1.8\LinearAlgebra\src\blas.jl:1514
[3] gemm_wrapper!(C::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, tA::Char, tB::Char, A::Matrix{Float32}, B::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool})
# LinearAlgebra C:\Users\Herr LU\AppData\Local\Programs\Julia-1.8.2\share\julia\stdlib\v1.8\LinearAlgebra\src\matmul.jl:674
[4] mul!
# C:\Users\Herr LU\AppData\Local\Programs\Julia-1.8.2\share\julia\stdlib\v1.8\LinearAlgebra\src\matmul.jl:161 [inlined]
[5] mul!
# C:\Users\Herr LU\AppData\Local\Programs\Julia-1.8.2\share\julia\stdlib\v1.8\LinearAlgebra\src\matmul.jl:276 [inlined]
[6] *
# C:\Users\Herr LU\AppData\Local\Programs\Julia-1.8.2\share\julia\stdlib\v1.8\LinearAlgebra\src\matmul.jl:148 [inlined]
[7] (::Dense{typeof(identity), Matrix{Float32}, Vector{Float32}})(x::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
# Flux C:\Users\Herr LU\.julia\packages\Flux\ZdbJr\src\layers\basic.jl:172
[8] (::MyDense)(x::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
# Main e:\Master Thesis\lu_jizhou\Learning\FluxTrial.jl:14
[9] macro expansion
# C:\Users\Herr LU\.julia\packages\Flux\ZdbJr\src\layers\basic.jl:53 [inlined]
[10] _applychain(layers::Tuple{Dense{typeof(tanh), CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(identity), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, MyDense, typeof(softmax)}, x::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
# Flux C:\Users\Herr LU\.julia\packages\Flux\ZdbJr\src\layers\basic.jl:53
[11] (::Chain{Tuple{Dense{typeof(tanh), CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(identity), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, MyDense, typeof(softmax)}})(x::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
# Flux C:\Users\Herr LU\.julia\packages\Flux\ZdbJr\src\layers\basic.jl:51
[12] top-level scope
# e:\Master Thesis\lu_jizhou\Learning\FluxTrial.jl:25
How can I solve the issue?