Hi,

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

C.sample()

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.