> 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.loss.utils

## Module Contents

### Functions

| Name                                                                                         | Description                                                                   |
| -------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
| [`_get_final_hidden_states`](#nemo_automodel-components-loss-utils-_get_final_hidden_states) | Return the final hidden-states tensor from an HF-like model output.           |
| [`_get_lm_head_module`](#nemo_automodel-components-loss-utils-_get_lm_head_module)           | Return the model's LM-head module, if one can be found.                       |
| [`_get_lm_head_weight`](#nemo_automodel-components-loss-utils-_get_lm_head_weight)           | Return the model's LM-head weight, materializing DTensor weights when needed. |
| [`calculate_loss`](#nemo_automodel-components-loss-utils-calculate_loss)                     | Calculate the loss.                                                           |

### API

```python
nemo_automodel.components.loss.utils._get_final_hidden_states(
    model_output: typing.Any
) -> typing.Optional[typing.Any]
```

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

Local copy of `components.training.model_output_utils.get_final_hidden_states`
to keep `components/loss/` import-independent from `components/training/`.

```python
nemo_automodel.components.loss.utils._get_lm_head_module(
    model: torch.nn.Module
) -> typing.Optional[torch.nn.Module]
```

Return the model's LM-head module, if one can be found.

Local copy of `components.utils.model_utils.get_lm_head_module` to keep
`components/loss/` import-independent from `components/utils/` (see the
`Components must not import each other` import-linter contract).

```python
nemo_automodel.components.loss.utils._get_lm_head_weight(
    model: torch.nn.Module
) -> torch.Tensor
```

Return the model's LM-head weight, materializing DTensor weights when needed.

```python
nemo_automodel.components.loss.utils.calculate_loss(
    loss_fn,
    kwargs = {}
) -> torch.Tensor
```

Calculate the loss.

**Parameters:**

Loss function.

Keyword arguments for the loss function.

**Returns:** `torch.Tensor`

The loss.