The last layer of my model is `self.linear = nn.Linear(4096, 2)`

. Usually to find out which class `data`

belongs to I write

```
output = model(data) #which return a tensor of size 2
pred = output.data.max(1, keepdim=True)[1]
```

Only I want the probability that `data`

belongs to the first or second class, so I thought I would write :

```
output = model(data) #which return a tensor of size 2
pred = torch.nn.Softmax(output) # pred = tensor([Prob(y=0|data), Prob(y=1|data)]) ?
```

But i does not work it gives me tensor of size two with negatives values. Do you have an idea to get this probabilities ?

Thanks