Pytorch: Softmax на каждом ряду 2-го Тензорного

Question

Мне нужна моя нейронная сеть для вывода N распределений по действиям A.каждое распределение должно проходить через softmax.Очевидно, что сумма каждой строки должна быть равна 1, а сумма всего слоя должна быть равна N. Есть ли такая функциональность в PyTorch?

Answer 1

Используйте softmax и укажите строку в качестве размера для операции

import torch.nn.Functional as F

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