Right now PyTorch batch_norm accepts any weight and running_mean shapes (eg 128 and 512) and doesn't complain at all, these checks https://github.com/torch/nn/blob/master/BatchNormalization.lua#L78 should probably be moved to C/CUDA side?