softmax_out is a (5, 932, 1682) torch tensor
m_hat = torch.stack([(r+1)*output for r, output in enumerate(softmax_out)], 0)
After this  m_hat is still a (5, 932, 1682) tensor. What does enumerate return? Why it can still keep its shape?