nemo_rl.utils.sequence_lengths#

Module Contents#

Functions#

to_cpu_int_tuple

Normalize sequence metadata at the host/device API boundary.

Data#

CpuIntTuple

API#

nemo_rl.utils.sequence_lengths.CpuIntTuple#

None

nemo_rl.utils.sequence_lengths.to_cpu_int_tuple(
values: torch.Tensor | collections.abc.Sequence[int],
) → nemo_rl.utils.sequence_lengths.CpuIntTuple#

Normalize sequence metadata at the host/device API boundary.

A CUDA tensor incurs one synchronization here. Callers should therefore invoke this before the model forward whenever the original CPU values are available. Code below this boundary accepts only :class:CpuIntTuple.