Migrate multinomial from the TH to ATen (CPU) (#25274)
Summary:
https://github.com/pytorch/pytorch/issues/24738
I updated the way to define n_categories and n_dist to fix https://github.com/pytorch/pytorch/issues/12309
Previously:
n_dist = prob_dist.size(0)
n_categories = prob_dist.size(1)
Changed to:
n_dist = prob_dist.size(-2)
n_categories = prob_dist.size(-1)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/25274
Differential Revision: D17137157
Pulled By: ifedan
fbshipit-source-id: 0320eafdaa7c272e169101b436b6c2ea4ba4736b
