hello, i have corrected the backward pass to train properly for the whole batch. Will you review it if i make a PR?