Improve documentation for torch.set_rng_state (#60422)

Summary:
Fixes https://github.com/pytorch/pytorch/issues/59974 by improving documentation for the function torch.set_rng_state

Pull Request resolved: https://github.com/pytorch/pytorch/pull/60422

Test Plan: Only a comment is being changed.

Reviewed By: bdhirsh

Differential Revision: D29281578

Pulled By: NivekT

fbshipit-source-id: 2c160f782438b7f91f16c44f06c342e8b8b8437b
This commit is contained in:
Kevin Tse
2021-06-22 07:09:24 -07:00
committed by Facebook GitHub Bot
parent 700df82881
commit 23bb2ed00a

View File

@ -8,6 +8,9 @@ import torch
def set_rng_state(new_state: torch.Tensor) -> None:
r"""Sets the random number generator state.
.. note: This function only works for CPU. For CUDA, please use
torch.manual_seed(seed), which works for both CPU and CUDA.
Args:
new_state (torch.ByteTensor): The desired state
"""