CODEWITHSUNDEEP
pythonpytorch
Granth
Granth

Reputation: 384

convert nn.Softmax to torch.tensor

I have a softmax function at the end of my Neural Network. I want to have the probabilities as torch.tensor.For it I am using torch.tensor(nn.softmax(x)) and getting error RuntimeError: Could not infer dtype of Softmax.

May I please know what I am doing wrong here or is there any other way to do it.

Upvotes: 2

Views: 2815

Answers (1)

Berriel
Berriel

Reputation: 13601

nn.Softmax is a class. You can use it like this:

import torch

x = torch.tensor([10., 3., 8.])
softmax = torch.nn.Softmax(dim=0)

probs = softmax(x)

or, you can use the Functional API torch.nn.functional.softmax:

import torch

x = torch.tensor([10., 3., 8.])

probs = torch.nn.functional.softmax(x, dim=0)

They are equivalent. In both cases, you can check that type(probs) is <class 'torch.Tensor'>.

Upvotes: 2

Related Questions