simplified batch softmax

Author: Taras Sereda

model.py

@@ -42,8 +42,7 @@ def __repr__(self):
 class InferenceBatchSoftmax(nn.Module):
     def forward(self, input_):
         if not self.training:
-            batch_size = input_.size()[0]
-            return torch.stack([F.softmax(input_[i], dim=1) for i in range(batch_size)], 0)
+            return F.softmax(input_, dim=-1)
         else:
             return input_