nemo_automodel.components.training.model_output_utils

Utilities for working with model output objects.

HuggingFace ModelOutput types typically store hidden_states as a tuple of tensors (tuple[Tensor, ...]) when output_hidden_states=True.

However, some custom models may store the final hidden state tensor directly in the hidden_states field (i.e., a single [B, T, H] tensor) to reduce memory/overhead.

Downstream training code should be robust to both representations.

Module Contents

Functions

Name	Description
`get_final_hidden_states`	Return the final hidden-states tensor from a HF-like model output.

Data

_TENSOR_LIKE

API

nemo_automodel.components.training.model_output_utils.get_final_hidden_states(
    model_output: typing.Any
) -> typing.Optional[typing.Any]

Return the final hidden-states tensor from a HF-like model output.

Supports both common layouts:

hidden_states is a tuple/list of tensors (HF default) → return last non-None entry
hidden_states is a single tensor-like object → return it as-is

Parameters:

model_output

Any

A HF ModelOutput-like object, or a plain dict with a hidden_states key.

Returns: Optional[Any]

The final hidden states tensor-like object, or None if not present.

nemo_automodel.components.training.model_output_utils._TENSOR_LIKE = (torch.Tensor, DTensor)