It seems that the lines with BS, C, H, W =x.shape are good only for channels_first and there is no code for BS, H, W, C. Also, what are the axes for softmax in those 2 cases?
It seems that the lines with BS, C, H, W =x.shape are good only for channels_first and there is no code for BS, H, W, C.
Also, what are the axes for softmax in those 2 cases?