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)
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 ?