# I am unable to fit a simple 2d function using a neural network in Julia. Am I doing somethng wrong here?

**URL:** <https://discourse.julialang.org/t/i-am-unable-to-fit-a-simple-2d-function-using-a-neural-network-in-julia-am-i-doing-somethng-wrong-here/103517>\
**Category:** General Usage\
**Tags:** flux, neural-network\
**Created:** [September 4, 2023, 7:35pm UTC](https://discourse.julialang.org/t/i-am-unable-to-fit-a-simple-2d-function-using-a-neural-network-in-julia-am-i-doing-somethng-wrong-here/103517 "2023-09-04T19:35:09Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![Shashank](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/shashank/32/12323_2.png) [@Shashank](https://discourse.julialang.org/u/Shashank)\
**Post date:** [September 4, 2023, 7:35pm UTC](https://discourse.julialang.org/t/i-am-unable-to-fit-a-simple-2d-function-using-a-neural-network-in-julia-am-i-doing-somethng-wrong-here/103517/1 "2023-09-04T19:35:09Z")

</div>

I am writing to write a code for interpolating a function using a neural network in Flux.jl, but the loss function does not decrease at all. Am I doing something wrong here? I am sorry if there is a trivial problem, as I am new to using Flux.jl. The same code works with slight modifications for a 1d function.

```julia
using Flux: train!
using Flux
using Statistics
x_train = rand(Float64, (2, 100))
y_train = x_train[1,:] + x_train[2,:]

predict = f64(Chain(Dense(2,64,leakyrelu),
                Dense(64,64,leakyrelu),
                Dense(64,1,leakyrelu)))
loss(model, x, y) = mean(abs2.(model(x) .- y))
data = [(x_train, y_train)]
opt = Flux.setup(Adam(0.01), predict)

for epoch in 1:5000
    train!(loss, predict, data, opt)
    println(loss(predict,x_train,y_train))
end
y_predict=predict(x_train)
for i in 1:100
    println(x_train[1,i]," ",x_train[2,i]," ",y_predict[i]," ",y_train[i])
end

```

---

<div class="post-metadata">

**Author:** ![skleinbo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/skleinbo/32/36080_2.png) [@skleinbo](https://discourse.julialang.org/u/skleinbo)\
**Post date:** [September 5, 2023, 7:51am UTC](https://discourse.julialang.org/t/i-am-unable-to-fit-a-simple-2d-function-using-a-neural-network-in-julia-am-i-doing-somethng-wrong-here/103517/2 "2023-09-05T07:51:13Z")

</div>

Replace

```julia
y_train = x_train[1,:] + x_train[2,:]

```

by

```julia
y_train = reshape(x_train[1,:] + x_train[2,:], 1, :)

```

to make the shape of `y_train` match that of `x_train`. The last dimension is the batch dimension.  
Also, use `mean(abs2, model(x).-y)` to avoid an unnecessary intermediate allocation.

In your original code, `predict(x_train).-y_train` produces

```julia
julia> predict(x_train).-y_train
100×100 Matrix{Float64}:
 -0.000221335 0.165909 0.562235 1.10108 0.422628 0.764869 … -0.0057506 0.900018 0.491736 0.215371 0.467521
 -0.16728 -0.00114932 0.395176 0.934021 0.25557 0.597811 -0.172809 0.732959 0.324677 0.0483125 0.300463
 -0.561125 -0.394994 0.00133184 0.540176 -0.138275 0.203966 -0.566654 0.339115 -0.0691673 -0.345532 -0.093382
 -1.10098 -0.934844 -0.538519 0.00032562 -0.678125 -0.335884 -1.1065 -0.200736 -0.609018 -0.885383 -0.633232
[...]

```

due to the broadcast `.` Clearly not what you intended!

---

<div class="post-metadata">

**Author:** ![Shashank](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/shashank/32/12323_2.png) [@Shashank](https://discourse.julialang.org/u/Shashank)\
**Post date:** [September 5, 2023, 12:01pm UTC](https://discourse.julialang.org/t/i-am-unable-to-fit-a-simple-2d-function-using-a-neural-network-in-julia-am-i-doing-somethng-wrong-here/103517/3 "2023-09-05T12:01:57Z")

</div>

Thanks a lot! That made the code work. I was working under the wrong assumption that I would get an error if there were a dimensions mismatch.

---

<div class="post-metadata">

**Author:** ![skleinbo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/skleinbo/32/36080_2.png) [@skleinbo](https://discourse.julialang.org/u/skleinbo)\
**Post date:** [September 5, 2023, 12:12pm UTC](https://discourse.julialang.org/t/i-am-unable-to-fit-a-simple-2d-function-using-a-neural-network-in-julia-am-i-doing-somethng-wrong-here/103517/4 "2023-09-05T12:12:25Z")

</div>

You do if you use `Flux.Losses.mse` in your loss function. That said, I agree that it would be prudent for a high-level function such as `train!` to give a warning if the input shapes are likely not as intended.

[https://fluxml.ai/Flux.jl/stable/models/losses/](https://fluxml.ai/Flux.jl/stable/models/losses/)