nemo_automodel.components.training.batch

Move model input tensors while retaining host-side packing metadata.

Module Contents

Functions

Name	Description
`move_batch_to_device`	Transfer training/validation inputs without moving deferred attention IDs.

API

nemo_automodel.components.training.batch.move_batch_to_device(
    batch: collections.abc.Mapping[str, typing.Any],
    device: torch.device
) -> dict[str, typing.Any]

Transfer training/validation inputs without moving deferred attention IDs.

Parameters:

batch

Mapping[str, Any]

Model inputs containing tensors of arbitrary shape, including nested mappings/lists/tuples. _packed_seq_ids is an integer tensor [batch, sequence] retained on CPU only when masks/Flash metadata still need preparation. THD and prepared-mask inputs transfer normally.

device

torch.device

Destination for all other tensor inputs.

Returns: dict[str, Any]

New containers preserving tensor shapes and dtypes. Deferred attention