Speed up subsample messenger
@@ -49,8 +49,7 @@ class _Subsample(Distribution):
if subsample_size is None or subsample_size >= self.size:
result = torch.arange(self.size, device=self.device)
else:
- result = torch.multinomial(torch.ones(self.size), self.subsample_size,
- replacement=False).to(self.device)
+ result = torch.randperm(self.size, device=self.device)[:subsample_size].clone()
return result.cuda() if self.use_cuda else result
def log_prob(self, x):
