Get the next state of a generator in Pytorch

August 9, 2023

I initialized a generator in PyTorch with a manual seed. How do I obtain the current state/seed of the generator? For example, if I want to check the state in a loop, how can I do that? This is what I want to achieve:

import torch

gen = torch.Generator("cpu").manual_seed(42)

for _ in range(5):
    # print the next state of the generator

PS: I tried to use gen.seed() but it gives different results on different runs.

>Solution :

You can use torch.Generator.get_state to retrieve the current state of the generator.

Per the docs:

get_state() → Tensor

Returns the Generator state as a torch.ByteTensor.

Returns: A torch.ByteTensor which contains all the necessary bits to restore a Generator to a specific point in time.

You can pass the value you get from get_state to set_state to restore the previous state of the generator.