Loss and output(model)

1

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

2

nn.Softmax is an nn.Module so you would either have to create the module first and call it with output or use F.softmax instead:

pred = torch.nn.Softmax(dim=1)(output)
# or
pred = F.softmax(output, dim=1)
1 Like