Any clues what might be causing this JIT issue? I would assume it should detect the Generator annotation. This is linked to following PR
Adding RNG to model initializers
Traceback (most recent call last):
File "/var/lib/jenkins/workspace/test/test_jit.py", line 14072, in test_nn_init
cu = torch.jit.CompilationUnit(code)
File "/opt/conda/lib/python3.6/site-packages/torch/jit/_recursive.py", line 838, in try_compile_fn
return torch.jit.script(fn, _rcb=rcb)
File "/opt/conda/lib/python3.6/site-packages/torch/jit/_script.py", line 1311, in script
qualified_name, ast, _rcb, get_default_args(obj)
RuntimeError:
Unknown type name 'torch.Generator':
File "/opt/conda/lib/python3.6/site-packages/torch/nn/init.py", line 122
def uniform_(tensor: Tensor, a: float = 0., b: float = 1., generator: torch.Generator = None) -> Tensor:
~~~~~~~~~~~~~~~ <--- HERE
r"""Fills the input Tensor with values drawn from the uniform
distribution :math:`\mathcal{U}(a, b)`.
'uniform_' is being compiled since it was called from 'test'
File "<string>", line 4
def test(a):
# type: (Tensor)
return torch.nn.init.uniform_(a)
~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE