There is an error when the no_atoms is different than the batch size.
I copied Huber Loss from here And now it is working when no_atoms different than batch size.