BatchNorm layer is broken if bottom blob has num_axes < 4

[kostinalexey]

This can happen when BatchNorm layer sits on top of FullyConnected layer. In this case input blob will have num_axes == 2.

Here is an example of error message
F0322 15:08:58.231948 19174 blob.hpp:251] Check failed: axis_index < num_axes() (2 vs. 2) axis 2 out of range for 2-D Blob with shape 1 512 (512)

The problem is in LayerSetUp function in both batch_norm_layer.cpp and cudnn_batch_norm_layer.cpp. At the end of the function there are following lines which causing the error:

  int N = bottom[0]->shape(0);
  int C = bottom[0]->shape(1);
  int H = bottom[0]->shape(2);
  int W = bottom[0]->shape(3);