[Discussion] fix zero divison error (Reformer batch size bug)
This PR is for discussion
During the training of the reformer model, I noticed that when you increase the batch size, a zero divison error often occurs
With these changes the error no longer occurs