def cuda_rng_load_state(rng_state):
    """
    Loads the state of cuda RNGs.

    :param rng_state: The list of RNG states (one per cuda device).
    """
    n_devices = torch.cuda.device_count()
    n_states = len(rng_state)
    if n_states < n_devices:
        warnings.warn(
            "Problem when reloading the state of torch.cuda RNGs: the given"
            "checkpoint contain a number of RNG states less than the "
            "number of currently available cuda devices "
            f"(got {n_states}, expected {n_devices}). "
            f"The RNG of cuda devices with ID >= {n_states} will not be "
            f"initialized!"
        )

    # The standard situation is n_devices == n_states: just re-load the state
    # of the RNGs of currently available GPUs. This reasoning also applies if
    # n_states > n_devices. However, this is a bad fit for n_states < n_devices
    # (and because of this, we show the warning above).
    for device_id, rng_state in enumerate(rng_state[:n_devices]):
        torch.cuda.set_rng_state(rng_state, f"cuda:{device_id}")

    Loads the state of cuda RNGs.

    :param rng_state: The list of RNG states (one per cuda device).
    
avalanche/training/determinism/cuda_rng.py