# Can you solve a NeuralODEMM on GPU?

**URL:** <https://discourse.julialang.org/t/can-you-solve-a-neuralodemm-on-gpu/57480>\
**Category:** Machine Learning\
**Tags:** diffeq, flux\
**Created:** [March 18, 2021, 3:03pm UTC](https://discourse.julialang.org/t/can-you-solve-a-neuralodemm-on-gpu/57480 "2021-03-18T15:03:01Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![bad\_at\_math](https://avatars.discourse-cdn.com/v4/letter/b/51bf81/32.png) [@bad\_at\_math](https://discourse.julialang.org/u/bad_at_math)\
**Post date:** [March 18, 2021, 3:03pm UTC](https://discourse.julialang.org/t/can-you-solve-a-neuralodemm-on-gpu/57480/1 "2021-03-18T15:03:01Z")

</div>

I’m trying to solve a neural ODE with a mass matrix using CUDA. Is this possible, or are the required solvers not GPU-compatible? I can define a `NeuralODEMM` with a `CuArray` Chain, but I get the error `ERROR: LoadError: scalar getindex is disallowed` when I try to apply the model to an initial condition.

Thank you!

MWE, based on here: [https://github.com/SciML/DiffEqFlux.jl/blob/master/test/neural\_ode\_mm.jl](https://github.com/SciML/DiffEqFlux.jl/blob/master/test/neural_ode_mm.jl) and here: [https://diffeqflux.sciml.ai/dev/GPUs/](https://diffeqflux.sciml.ai/dev/GPUs/)

```julia
using Flux, DiffEqFlux, GalacticOptim, OrdinaryDiffEq
using CUDA; CUDA.allowscalar(false)

# generate training data
function f(du,u,p,t)
    y₁,y₂,y₃ = u
    k₁,k₂,k₃ = p
    du[1] = -k₁*y₁ + k₃*y₂*y₃
    du[2] = k₁*y₁ - k₃*y₂*y₃ - k₂*y₂^2
    du[3] = y₁ + y₂ + y₃ - 1
    nothing
end

# initial condition
u₀ = [1.0, 0, 0]
# mass matrix
M = [1. 0 0
     0 1. 0
     0 0 0]
tspan = (0.0,1.0)
p = [0.04, 3e7, 1e4]
func = ODEFunction(f, mass_matrix=M)
prob = ODEProblem(func, u₀, tspan, p)
sol = solve(prob, Rodas5(), saveat=0.1)

# define model
dudt2 = Chain(Dense(3, 64, tanh), Dense(64, 2)) |> gpu
ndae = NeuralODEMM(dudt2,
                   (u,p,t) -> [u[1] + u[2] + u[3] - 1],
                   tspan,
                   M,
                   Rodas5(autodiff=false),
                   saveat=0.1)
u0_gpu = u₀ |> gpu
display(typeof(ndae.p)) # shows CuArray{Float32,1}
display(typeof(ndae(u0_gpu))) # throws error ERROR: LoadError: scalar getindex is disallowed

```