Commit 977e36e5 authored by Myle Ott's avatar Myle Ott Committed by Facebook Github Bot
Browse files

Support multiple seeds in data_utils.numpy_seed

Summary: Pull Request resolved: https://github.com/fairinternal/fairseq-py/pull/610

Differential Revision: D15541261

Pulled By: myleott

fbshipit-source-id: f0b823cf4f04c5ef3205f6d259c6dcad4cc329b1
parent 3e472b22
......@@ -44,12 +44,14 @@ def collate_tokens(values, pad_idx, eos_idx=None, left_pad=False, move_eos_to_be
@contextlib.contextmanager
def numpy_seed(seed):
def numpy_seed(seed, *addl_seeds):
"""Context manager which seeds the NumPy PRNG with the specified seed and
restores the state afterward"""
if seed is None:
yield
return
if len(addl_seeds) > 0:
seed = int(hash((seed, *addl_seeds)) % 1e6)
state = np.random.get_state()
np.random.seed(seed)
try:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment