Suspicious `sum` in nn model training

[vdyma]

There is a suspicious sum in the nn model training.