Pytorch: Softmax on each row of a 2d Tensor
I need my neural net to output N distributions over A actions. each distribution should go through softmax. The sum of each row should then obviously be 1 and the sum of the whole layer should be N. Is there such functionality in PyTorch?

Use the softmax and specify the row as the dimension to operate on

import torch.nn.Functional as F

x = ...# your N x A input
x_distribution = F.softmax(x, dim = 1)

