Update docs/src/examples/ode/second_order_neural.md

SciML · Sep 28, 2024 · f5398e3 · f5398e3
1 parent dd6c89d
commit f5398e3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/docs/src/examples/ode/second_order_neural.md b/docs/src/examples/ode/second_order_neural.md
@@ -33,7 +33,7 @@ t = range(tspan[1], tspan[2], length = 20)
 model = Chain(Dense(2, 50, tanh), Dense(50, 2))
 ps, st = Lux.setup(Random.default_rng(), model)
 ps = ComponentArray(ps)
-model = Lux.StatefulLuxLayer(model, ps, st)
+model = Lux.StatefulLuxLayer{true}(model, ps, st)
 
 ff(du, u, p, t) = model(u, p)
 prob = SecondOrderODEProblem{false}(ff, du0, u0, tspan, ps)