> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# 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`](#nemo_automodel-components-training-model_output_utils-get_final_hidden_states) | Return the final hidden-states tensor from a HF-like model output. |

### Data

[`_TENSOR_LIKE`](#nemo_automodel-components-training-model_output_utils-_TENSOR_LIKE)

### API

```python
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:**

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.

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