* Support case where BN module has no learnable parameters (affine == False) * Support conv1d and conv3d