I see, it was my mistake. Haven’t checked the sizes properly.
Try to change the calculation to this line:
x.data = k.view(1, -1, 1, 1).expand_as(x) * x.data + beta.view(1, -1, 1, 1).expand_as(x).data

I’m sure it’s not the optimal way to calculate it, so maybe someone can propose a better way.
@SimonW Yeah you are right! The proposed calculations in the first post should work however, even though BN is nonlinear, or am I missing something?