Currently nn.Softmax() can only be applied on 2D tensor. Is that possible we can specify dimension? E.g.,

x = torch.FloatTensor(128, 5, 27)
softmax = nn.Softmax(axis=1)
y = softmax(Variable(x))
y.size()

y.size() == x.size() but y's second dimension (5) has been applied softmax operation.