Say I have a probability tensor T, which is shape (Batch, N, D) and is normalized along the last dimension. Ideally, I want to sample along this last dimension such that I can index into a variable with the same shape as T.
I have been trying to do this using pyro.distributions.Categorical to no avail. Would anyone be willing to point out what I'm currently doing wrong?
T = torch.randn(128,4,5)
T = ptsoftmax(T,-1) --- This is just a vanilla softmax function. Double checked that it works as intended.
C = pyro.distributions.Categorical(ps=torch.autograd.Variable(T))
This results in:
The expanded size of the tensor (4) must match the existing size (128) at non-singleton dimension 1. at /pytorch/torch/lib/TH/generic/THTensor.c:308
I have also tried creating C such that C = pyro.distributions.Categorical(ps=torch.autograd.Variable(T), vs=torch.autograd.Variable(T) )
to the same success.
Thank you very much.