How to reshape without sample function?

How do I reshape a Categorical distribution of dimension [800,3] to dimension [32,1,5,5] without using sample function? Because when i sample and then reshape it becomes a constant(because of .sample()) and I am unable to provide it to my neural network (with convolutional filter dimensions [32,1,5,5] for updating the distribution parameters.

Hmm, it would probably be easier to pass your logits tensor to a neural net, rather than the Categorical(logits=logits) distribution. I’d recommend simply working with logits until you need to either sample or compute log_prob().

@fritzo. Thank you for your reply. If I just pass the logits to the neural network and not use Categorical(logits=logits),then how to I create a variational posterior and optimize its parameters by maximizing the ELBO

Could you post some code to help me understand your problem? I’m having a tough time understanding how you intend to pass a categorical distribution through a neural net.