meipuru344:

grad = grad_fn(output, teacher[i]) #calc non scalar grad


You can divide the grad with the length of batch. It is needed especially when the batches are of different lengths.
If the batches are of same length, dividing by length of batch size can be thought to be absorbed by learning rate anyway, in which case its not needed.