Is my loss function wrong?

This would work if all batches contain the same number of samples. In case the last one has fewer samples, you would add a small error to the loss calculation as also described here.

1 Like