Any easy approach to apply nn.Softmax() along each dimension?

[yuandong-tian]

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.

[apaszke]

Right now there's no other way than transposing and flattening all other dimensions to a single one at the front (and reversing that after softmax). It's on the roadmap and is already tracked in #1020.