Solved by rahulvigneswaran in post #2 
              
              
                In case 1: model(state_batch) 

The output dimension of torch.nn.Linear is (3,2) - 2 Dimensional

In case 2: [model(s).detach().numpy() for s in state_batch] 

The output dimension of torch.nn.Linear is just (2) - 1 Dimensional

The torch.nn.Softmax(dim=0) does softmax at dim=0 which will work prope…