you should put a LogSoftmax in the policy network’s last layer, not Softmax

If I do so, I will be  required to  subsequently call exp(log_probs) to match a signature of the torch.distributions.Categorical.
Also, the subsequent call to torch.distributions.Categorica(probs).log_prob(0) will also take the log, which, again, will cause all the numerical issues.
So,  either I don’t understand the purpose of the current implementation of torch.distributions or the example on the github release page is misleading.
Also, probably something is weird with sampling from Categorical. After executing
probs = F.softmax(torch.autograd.Variable(torch.Tensor([.25, .6])))
dist = torch.distributions.Categorical(probs)

this code works find
torch.multinomial(dist.probs, 2, True)

but this one does not
dist.sample_n(10)

yielding the following stack trace
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-8-792c4214dc6f> in <module>()
----> 1 dist.sample_n(10)

~/anaconda3/lib/python3.6/site-packages/torch/distributions.py in sample_n(self, n)
    143             return self.sample().expand(1, 1)
    144         else:
--> 145             return torch.multinomial(self.probs, n, True).t()
    146 
    147     def log_prob(self, value):

RuntimeError: invalid argument 2: out of range at /Users/soumith/minicondabuild3/conda-bld/pytorch_1512381214802/work/torch/lib/TH/generic/THTensor.c:479

This happens because of the last .t().