oneflow.cuda.set_rng_state

oneflow.cuda.set_rng_state(new_state: oneflow.Tensor, device: Union[int, str, oneflow._oneflow_internal.device] = 'cuda') → None

Sets the random number generator state of the specified GPU.

Parameters

• new_state (flow.ByteTensor) – The desired state

• device (flow.device or int, optional) – The device to set the RNG state. Default: 'cuda' (i.e., flow.device('cuda'), the current CUDA device).