Thanks, Siman!
That makes a lot of sense now.
Resolved it by changing a line in the backward pass in AdBias to
if len_grad == 2: grad_bias = -1 * torch.mean(grad_out)
The example above was only meant to illustrate my issue with broadcasting doing the learning phase. Thanks for the suggestion though.