torch.distributions.categorical.Categorical not producing tensor output of same dimension

1

Hi,

I have a 3 dimension tensor “prob: tensor([[[0.2793, 0.3314, 0.3893]]], grad_fn=)”
and I am sampling an action by Categorical distribution.

                    print("prob:",prob)
                    action = Categorical(prob).sample().detach()
                    print("action:",action)
                    log_prob_a = log_prob.gather(1, action)

This is the output I am getting.

prob: tensor([[[0.3201, 0.3268, 0.3531]]], grad_fn=)
action: tensor([[2]])
Traceback (most recent call last):
File “D:\get_data\ind_a3c_lstm.py”, line 99, in train
log_prob_a = log_prob.gather(1, action)
RuntimeError: Index tensor must have the same number of dimensions as input tensor

My, Question is why action is not have same number of dimensions as prob tensor.