Error in solution representation when using `Lux.BatchNorm`

[NoFishLikeIan]

Evaluating the solution of an optimization procedure using

discretization = PhysicsInformedNN(chain, QuadratureTraining())
prob = discretize(system, discretization)

result = Optimization.solve(prob, BFGS())

discretization.phi([0.5], result.u) # Error here

yields the error

BoundsError: attempt to access Tuple{Int64} at index [0]

if the Lux.chain contains a Lux.BatchNorm layer.

Here is a MWE to reproduce the issue,

using Lux: Chain, Dense, BatchNorm, relu
using NeuralPDE
using Optimization, OptimizationOptimJL
using ModelingToolkit: Interval, infimum, supremum

n = 18
chain = Chain(
    Dense(1, n),
    BatchNorm(n, relu), # Without this line there is no error
    Dense(n, 1)
)

@parameters t
@variables f(..)

D = Differential(t)

diffeq = [ f(t) ~ D(f(t)) ]
bcs = [ f(0) ~ 1 ]

tdomain = t ∈ Interval(0, 1)

@named system = PDESystem(diffeq, bcs, [tdomain], [t], [f(t)])

discretization = PhysicsInformedNN(chain, QuadratureTraining())
prob = discretize(system, discretization)

result = Optimization.solve(prob, BFGS())

discretization.phi([0.5], result.u) # Error occurs here

The Pkg.status() yields

⌅ [b2108857] Lux v0.4.58
  [961ee093] ModelingToolkit v8.63.0
  [315f7962] NeuralPDE v5.7.0
  [7f7a1694] Optimization v3.15.2
  [36348300] OptimizationOptimJL v0.1.9

A full Stacktrace can be found here.