Was trying to train a PINN for NLSE using NeuralPDE. Everything runs fine on CPU, but when I try to move calculations to GPU, I get the following error:

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

My code is below. The error is raised at `prob=discretize...`

line.

```
using NeuralPDE, Lux, LuxCUDA, Random, ComponentArrays
using Optimization
using OptimizationOptimisers
import ModelingToolkit: Interval
@parameters t x
@variables u(..)
Dt = Differential(t)
Dxx = Differential(x)^2
L = 10.0
w = 0.2
A = 1/w
ic(x) = A / cosh(x/w)^2
eq = Dt(u(t, x)) ~ 0.5 * Dxx(u(t, x)) + u(t, x) * abs(u(t, x))^2
bcs = [u(t, 0) ~ u(t, L), u(0, x) ~ ic(x)]
domains = [t ∈ IntervalDomain(0.0, 10.0), x ∈ IntervalDomain(-L/2, L/2)]
chain = Lux.Chain(Lux.Dense(2, 16, Lux.σ), Lux.Dense(16, 16, Lux.σ), Lux.Dense(16, 1))
strategy = GridTraining(0.05)
ps_cpu = Lux.setup(Random.default_rng(), chain)[1]
ps = ps_cpu |> ComponentArray |> Lux.gpu_device() .|> Float64
discretization = PhysicsInformedNN(chain, strategy, init_params=ps)
@named pde_system = PDESystem(eq, bcs, domains, [t, x], [u(t,x)])
prob = discretize(pde_system, discretization)
callback = function (p, l)
#println("Current loss is: $l")
return false
end
res = Optimization.solve(prob, Adam(0.01); callback = callback, maxiters = 2500)
```

**EDIT:** Updated the code to my last version. Also ignore the equation being wrong (NLSE has imaginary unit multiplied by time derivative, mine doesn’t)