nemo_rl.utils.tensor_ops#

Module Contents#

Functions#

pad_and_concat

Right-pad each tensor along dim 1 to target_len, then concatenate along dim 0.

API#

nemo_rl.utils.tensor_ops.pad_and_concat(
tensors: list[torch.Tensor],
*,
target_len: int,
pad_value: float = 0.0,
) → torch.Tensor#

Right-pad each tensor along dim 1 to target_len, then concatenate along dim 0.

Parameters:

  • tensors – Tensors with the sequence on dim 1.

  • target_len – Sequence length every tensor is padded up to.

  • pad_value – Fill value for the padded positions.

Returns:

The padded tensors concatenated along dim 0.