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#

get_final_hidden_states

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

API#

nemo_automodel.components.training.model_output_utils.get_final_hidden_states(
model_output: Any,
) → Optional[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 – A HF ModelOutput-like object, or a plain dict with a hidden_states key.

Returns:

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