As you’ve already described the number of channels in the input activation to the batchnorm layer doesn’t match the expected channels, so you would have to permute the activation via:
x = x.permute(0, 2, 1)

before passing it to layer_b.