BoundsError during training process

[jcvillaquira]

Hi, I am training a NeuralOperator to model a map that takes two one-dimensional functions to a two-dimensional function. However it always gives error somewhere in the training process. Here is a minimal working example

# nt = temporal discretisation, nx = spatiall discretisation
nt = 10
nx = 20
n_samples = 100

# Generate data.
x_data = rand(2, nx, n_samples) # Two functions as input
y_data = rand(nt, nx, n_samples) # Two-dimensional function with size nt x nx
data_train, data_test = splitobs((Float32.(x_data), Float32.(y_data)), at = 0.9)
loader = DataLoader(data_train), DataLoader(data_test)

# Create model.
model = FourierNeuralOperator(ch = (2, 64, 64, 64, 64, 64, 128, nt), σ = gelu)
optimiser = Flux.Optimiser(WeightDecay(1.0f-4), Flux.Adam(1.0f-3))
data = collect.(loader)

# Train model
learner = Learner(model, data, optimiser, l₂loss)
fit!(learner, 1)

I always get BoundsError in the code

ERROR: BoundsError: attempt to access 11×64×1 Array{ComplexF32, 3} at index [1:16, 1:64, 1:1]

Thanks!