Глубокое Q-Learning: сбой torch.nn.functional.softmax

Question

Я следую учебному пособию, и функция softmax падает, когда я его использую.

newSignals = [0.5, 0., 0., -0.7911, 0.7911]
newState = torch.Tensor(newSignals).float().unsqueeze(0)
probs = F.softmax(self.model(newState), dim=1)

self.model - это нейронная сеть (torch.nn.module), которая возвращает тензор, подобный

tensor([[ 0.2699, -0.2176, 0.0333]], grad_fn=<AddmmBackward>)

Итак, строка probs = F.softmax(self.model(newState), dim=1) завершает работу программы, но когда dim=0 работает, но это не хорошо.

Answer 1

Отказ от ответственности: я сожалею, что это, вероятно, должен был быть комментарий, но я не могу написать все ниже в комментарии.

Вы уверены, что это проблема?Ниже фрагмент просто работал для меня.

import torch
a = torch.tensor([[ 0.2699, -0.2176,  0.0333]]) 
a.softmax(dim=1)
> tensor([[0.4161, 0.2555, 0.3284]])
a.softmax(dim=0)
> tensor([[1., 1., 1.]])