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


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.

                    action = Categorical(prob).sample().detach()
                    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\”, 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.