Here is what BN does:
assumes that inputs has shape:  [B x C x *]
  B is minibatch size
  C is the size of second dimension. In image data, it is the channel 
    dim. But generally it is the second dim.
  * is arbitrary number of dimensions of arbitrary sizes

Normalizes each slice input[:, i, :] for i = 0, ..., C - 1, that is
normalizing the slice using the mean and std computed from that slice.

So if you have input of shape [B, *], and do bn_layer(input.unsqueeze(0)), you are effectively doing:
input_new = input.unsqueeze(0) of shape [1, B, *]
output = bn(input_new) normalizes each input_new[:, i] = input[i] slice
  for i = 0, ..., B - 1.
output = output.squeeze(0) gives what you want