Add utils.with_torch_seed
Summary: Pull Request resolved:
@@ -406,6 +406,17 @@ def set_torch_seed(seed):
torch.cuda.manual_seed(seed)
+@contextlib.contextmanager
+def with_torch_seed(seed):
+ assert isinstance(seed, int)
+ rng_state = torch.get_rng_state()
+ cuda_rng_state = torch.cuda.get_rng_state()
+ set_torch_seed(seed)
+ yield
+ torch.set_rng_state(rng_state)
+ torch.cuda.set_rng_state(cuda_rng_state)
+
+
def parse_alignment(line):
"""
Parses a single line from the alingment file.
