Type check quasirandom (pytorch#45434)

Summary:
Fixes pytorch#42978.

Pull Request resolved: pytorch#45434

Reviewed By: walterddr

Differential Revision: D23967139

Pulled By: ajitmaths

fbshipit-source-id: bcee6627f367fd01aa9a5c10a7c24331fc1823ad

mypy.ini

@@ -90,9 +90,6 @@ ignore_errors = True
 [mypy-torch.quantization.fx.*]
 ignore_errors = True
 
-[mypy-torch.quasirandom]
-ignore_errors = True
-
 [mypy-torch.distributions.*]
 ignore_errors = True
 

torch/quasirandom.py

@@ -1,4 +1,5 @@
 import torch
+from typing import Optional
 
 
 class SobolEngine(object):
@@ -57,11 +58,10 @@ def __init__(self, dimension, scramble=False, seed=None):
         torch._sobol_engine_initialize_state_(self.sobolstate, self.dimension)
 
         if self.scramble:
+            g: Optional[torch.Generator] = None
             if self.seed is not None:
                 g = torch.Generator()
                 g.manual_seed(self.seed)
-            else:
-                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)))