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