Summary:
Fixes https://github.com/pytorch/pytorch/issues/63723
Since it's a generator function the type annotation shall be `Generator`.
![image](https://user-images.githubusercontent.com/
47299190/
130318830-
29ef9529-0daa-463c-90b2-
1b11f63ade8a.png)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63724
Reviewed By: iramazanli
Differential Revision:
D30543098
Pulled By: heitorschueroff
fbshipit-source-id:
ebdd34749defe1e26c899146786a0357ab4b4b9b
import contextlib
+from typing import Generator
import warnings
from torch._C import default_generator
@contextlib.contextmanager
-def fork_rng(devices=None, enabled=True, _caller="fork_rng", _devices_kw="devices"):
+def fork_rng(devices=None, enabled=True, _caller="fork_rng", _devices_kw="devices") -> Generator:
"""
Forks the RNG, so that when you return, the RNG is reset
to the state that it was previously in.