Thanks for the solution @smth!
If I can change the per-channel operation to one element-wise tensor multiplication by expanding the mean and std in proper way, then is it gonna be faster? Or since its graph, it will not affect much?
Something like below,
t_mean = torch.FloatTensor(mean).view(3,1,1).expand(3,H,W).contiguous().view(1,3,H,W)
t_std = torch.FloatTensor(std).view(3,1,1).expand(3,H,W).contiguous().view(1,3,H,W)
x = y * t_std.expand(B,3,H,W) +t_mean.expand(B,3,H,W)

I’m not sure what is most simple way to change the shape properly though…