# Why doesn't the loss calculated by Flux \`withgradient\` match what I have calculated?

**URL:** <https://discourse.julialang.org/t/why-doesnt-the-loss-calculated-by-flux-withgradient-match-what-i-have-calculated/109228>\
**Category:** Machine Learning\
**Tags:** question, flux\
**Created:** [January 25, 2024, 12:36am UTC](https://discourse.julialang.org/t/why-doesnt-the-loss-calculated-by-flux-withgradient-match-what-i-have-calculated/109228 "2024-01-25T00:36:13Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![kirklong](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/kirklong/32/19433_2.png) [@kirklong](https://discourse.julialang.org/u/kirklong)\
**Post date:** [January 25, 2024, 12:36am UTC](https://discourse.julialang.org/t/why-doesnt-the-loss-calculated-by-flux-withgradient-match-what-i-have-calculated/109228/1 "2024-01-25T00:36:13Z")

</div>

Just posted this to [StackOverflow](https://stackoverflow.com/questions/77876955/why-doesnt-the-loss-calculated-by-flux-withgradient-match-what-i-have-calcula), but crossposting here for Julia reach.

TL;DR – when making a simple CNN I noticed weird behavior between the reported loss during training and the loss I actually calculated from the “trained” model (trained model looked like garbage despite loss saying it should be quite good). The crux of the problem I identified is that when having a `BatchNorm` layer, the loss value calculated appears to mutate after calling `withgradient`.

For example:

```julia
using Flux
x = rand(100,100,1,1) #say a greyscale image 100x100 with 1 channel (greyscale) and 1 batch
y = @. 5*x + 3 #output image, some relationship to the input values (doesn't matter for this)
m = Chain(BatchNorm(1),Conv((1,1),1=>1)) #very simple model (doesn't really do anything but illustrates the problem)
l_init = Flux.mse(m(x),y) #initial loss after model creation
l_grad, grad = Flux.withgradient(m -> Flux.mse(m(x),y), m) #loss calculated by gradient
l_final = Flux.mse(m(x),y) #loss calculated again using the model (no parameters have been updated)
println("initial loss: $l_init")
println("loss calculated in withgradient: $l_grad")
println("final loss: $l_final")

```

produces different values for each of the loss calculations (sometimes very different), whereas running a simple model without a `BatchNorm` layer:

```julia
using Flux
x = rand(100,100,1,1) #say a greyscale image 100x100 with 1 channel (greyscale) and 1 batch
y = @. 5*x + 3 #output image
m = Chain(Conv((1,1),1=>1))
l_init = Flux.mse(m(x),y) #initial loss after model creation
l_grad, grad = Flux.withgradient(m -> Flux.mse(m(x),y), m)
l_final = Flux.mse(m(x),y)
println("initial loss: $l_init")
println("loss calculated in withgradient: $l_grad")
println("final loss: $l_final")

```

produces the same output (which is what I expected).

Why does adding `BatchNorm` do this, and is this intended? Am I using it wrong?

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [January 26, 2024, 2:59pm UTC](https://discourse.julialang.org/t/why-doesnt-the-loss-calculated-by-flux-withgradient-match-what-i-have-calculated/109228/2 "2024-01-26T14:59:42Z")

</div>

I presume you’re used to PyTorch and having to manually specify whether a layer is in train or test/eval mode? Flux does this automatically by default, so taking gradients of a model will let BatchNorm detect that it should be updating running stats. If you’d like more manual control, `trainmode!` and `testmode!` are available per [the docs](https://fluxml.ai/Flux.jl/stable/models/layers/#Test-vs.-Train).

---

<div class="post-metadata">

**Author:** ![kirklong](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/kirklong/32/19433_2.png) [@kirklong](https://discourse.julialang.org/u/kirklong)\
**Post date:** [January 26, 2024, 9:05pm UTC](https://discourse.julialang.org/t/why-doesnt-the-loss-calculated-by-flux-withgradient-match-what-i-have-calculated/109228/3 "2024-01-26T21:05:57Z")

</div>

thanks, yes applying those explicitly makes the values the same!