instead of using mean(), I used sum and divided by batch_size, now I get correct results. I don’t understand why.