CuArray error when running NeuralPDE.jl on JuliaHub

Specific Domains GPU
,
1

Hi everyone!
I was trying to test the GPU ability of NeuralPDE.jl on the JuliaHub, my network is defined like this:

# Neural network
inner = 25
chain = Chain(Dense(3,inner,Lux.σ),
              Dense(inner,inner,Lux.σ),
              Dense(inner,inner,Lux.σ),
              Dense(inner,inner,Lux.σ),
              Dense(inner,6)) 

strategy = GridTraining(0.05)
ps = Lux.setup(Random.default_rng(), chain)[1]
ps = ps |> ComponentArray |> gpu .|> Float64
discretization = PhysicsInformedNN(chain,
                                   strategy,
                                   init_params = ps)

@named pdesystem = PDESystem(eqs, ic_bc, domains, [t,x,z], [u(t,x,z), w(t,x,z), P(t,x,z),
T(t,x,z), τ₁(t,x,z), τ₃(t,x,z) ])

prob = discretize(pdesystem,discretization)

But when I execute the last line I get:

ERROR: CuArray only supports element types that are allocated inline.
Real is not allocated inline

I was looking to see if this issue was mentioned somewhere and I found this

But I didn’t understand if this was solved yet.

This is my first time trying to use a GPU, so any hints would be highly appreciated!