深度Q-Learning：torch.nn.functional.softmax崩溃

Question

I am following a tutorial, and the function softmax crashes when I use it.

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 is a neural network ( torch.nn.module ), which return a Tensor like

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

So, the line probs = F.softmax(self.model(newState), dim=1) crash the program but when dim=0 it works but it is not good.

Answer 1

Disclaimer: I am sorry this probably should have been a comment but I can't write all the below in a comment.

Are you sure this is the problem? Below snippet just worked for me.

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.]])