categorical.py 文件源码

def sample(self, sample_shape=torch.Size()):
        num_events = self.probs.size()[-1]
        sample_shape = self._extended_shape(sample_shape)
        param_shape = sample_shape + self.probs.size()[-1:]
        probs = self.probs.expand(param_shape)
        probs_2d = probs.contiguous().view(-1, num_events)
        sample_2d = torch.multinomial(probs_2d, 1, True)
        return sample_2d.contiguous().view(sample_shape)