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 properly for Case 2 as there is only 1 dimension but in case of Case 1 there are 2 dimensions and that is the root cause for the issue.

Changing your code to torch.nn.Softmax(dim=1) should fix this issue for Case 1.

And if you still want the new changed code to work for Case 2, then do the following,
pred_batch = [model(s.unsqueeze(dim=0).detach().numpy() for s in state_batch]