How to do channel-wise scaling

I have a 4D tensor X in NxCxWxH like the output from Conv2d. Now, I like to scale up each channel with different constants, like

for i, each batch:
    for j, each channel:
       Y[i][j] = s[j]*W*H where s[j] is a scalar value 

I was able to write a for-loop to go over each image in a batch, with a bunch of view changes. I wonder if there is a vectorized way of doing per-channel constant scaling over batches.

Not sure why you are multiplying each s[j] with W*H is not not same as S = S * W * H outside the loop.

S = S * W * H  # assuming S.shape() = [C]
S = S.rehape(C, 1, 1)
Y = Y * S