[pytorch] Make behavior of SobolEngine consistent w/ other RNG functions (#36427)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/36427

Addresses https://github.com/pytorch/pytorch/issues/36341

Test Plan: unit tests

Reviewed By: ldworkin

Differential Revision: D20952703

fbshipit-source-id: 28055f4c4c0f8012c2d96e473b822fa455dd833c
This commit is contained in:
Max Balandat
2020-04-13 07:48:05 -07:00
committed by Facebook GitHub Bot
parent d2e0c628e9
commit 379e4d9cad
2 changed files with 14 additions and 2 deletions

View File

@ -57,11 +57,11 @@ class SobolEngine(object):
torch._sobol_engine_initialize_state_(self.sobolstate, self.dimension)
if self.scramble:
g = torch.Generator()
if self.seed is not None:
g = torch.Generator()
g.manual_seed(self.seed)
else:
g.seed()
g = None
shift_ints = torch.randint(2, (self.dimension, self.MAXBIT), device=cpu, generator=g)
self.shift = torch.mv(shift_ints, torch.pow(2, torch.arange(0, self.MAXBIT, device=cpu)))