Error when running a simple ANN model in Flux

Specific Domains Machine Learning
1

Hi I am trying to run a simple ANN model in flux but having trouble with understanding the errors. Any idea what I am doing incorrectly

# Shapes : X_train : (204,2) , X_val : (51,2), y_train : (204,3), y_val : (51,3) 
# Shapes of X_train, y_train transposed during training

in_dims = 2
l1 = 16
l2 = 16
out_dims = 3
lr = 0.001
epochs = 50

# Function to get Model
get_model(in_dims, l1, l2, out_dims) = Chain(
    Dense(in_dims, l1, relu),
    Dense(l1, l2, relu),
    Dense(l2, out_dims)
)

# Function for training
function training_loop(model::Chain,dataset::Dict)
    (;X_train,X_val,y_train,y_val) = (;sort(dataset)...)
    
    loss_fn(x,y) = mse(model(x),y)
    opt = ADAM(lr)

    losses = Float64[]
    for e = 1:epochs
        train!(loss_fn, params(model), [(X_train',y_train')], opt)
        train_loss = loss_fn(X_train, y_train)
        push!(losses, train_loss)
        println("epoch : $(e) \t training loss : $(train_loss)")
    end
    return losses
end

# Call the function to contruct model and train
dataset = Dict(:X_train => X_train, :X_val => X_val, :y_train => y_train, :y_val => y_val)
train_loss = training_loop(model, dataset)

>>>
ERROR: Mutating arrays is not supported -- called setindex!(Matrix{Any}, ...)
This error occurs when you ask Zygote to differentiate operations that change
the elements of arrays in place (e.g. setting values with x .= ...)
2

If you print out the sizes and eltypes of x and y in loss_fn, that would help identify any issues.