> 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.intermediate_distill

## Module Contents

### Classes

| Name                                                                                                      | Description                                                            |
| --------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------- |
| [`IntermediateDistillLoss`](#nemo_automodel-components-loss-intermediate_distill-IntermediateDistillLoss) | Intermediate-layer feature distillation module.                        |
| [`LayerCapture`](#nemo_automodel-components-loss-intermediate_distill-LayerCapture)                       | Forward-hook helper for capturing selected intermediate hidden states. |

### Functions

| Name                                                                                                            | Description                                                            |
| --------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------- |
| [`_lookup_layer`](#nemo_automodel-components-loss-intermediate_distill-_lookup_layer)                           | -                                                                      |
| [`_masked_per_token_loss`](#nemo_automodel-components-loss-intermediate_distill-_masked_per_token_loss)         | -                                                                      |
| [`_pick_projector`](#nemo_automodel-components-loss-intermediate_distill-_pick_projector)                       | -                                                                      |
| [`intermediate_loss_function`](#nemo_automodel-components-loss-intermediate_distill-intermediate_loss_function) | Per-token feature distillation between matched student/teacher layers. |
| [`intermediate_loss_pair`](#nemo_automodel-components-loss-intermediate_distill-intermediate_loss_pair)         | Sum of intermediate loss on query and doc sides.                       |

### Data

[`HiddenStatesLike`](#nemo_automodel-components-loss-intermediate_distill-HiddenStatesLike)

[`ProjectorLike`](#nemo_automodel-components-loss-intermediate_distill-ProjectorLike)

### API

```python
class nemo_automodel.components.loss.intermediate_distill.IntermediateDistillLoss(
    layer_pairs: collections.abc.Sequence[tuple[int, int]] | collections.abc.Sequence[list[int]],
    loss_type: str = 'mse',
    normalize: bool = False,
    layer_weights: collections.abc.Sequence[float] | None = None,
    reduction: str = 'mean'
)
```

**Bases:** `Module`

Intermediate-layer feature distillation module.

```python
nemo_automodel.components.loss.intermediate_distill.IntermediateDistillLoss.forward(
    s_q_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    t_q_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    s_d_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    t_d_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    attn_q: torch.Tensor | None = None,
    attn_d: torch.Tensor | None = None,
    projector: nemo_automodel.components.loss.intermediate_distill.ProjectorLike | None = None
) -> torch.Tensor
```

```python
class nemo_automodel.components.loss.intermediate_distill.LayerCapture(
    detach: bool = False
)
```

Forward-hook helper for capturing selected intermediate hidden states.

```python
nemo_automodel.components.loss.intermediate_distill.LayerCapture._make_hook(
    idx: int
)
```

```python
nemo_automodel.components.loss.intermediate_distill.LayerCapture.attach(
    layers: torch.nn.ModuleList,
    indices: collections.abc.Iterable[int]
) -> None
```

```python
nemo_automodel.components.loss.intermediate_distill.LayerCapture.detach_hooks() -> None
```

```python
nemo_automodel.components.loss.intermediate_distill.LayerCapture.reset() -> None
```

```python
nemo_automodel.components.loss.intermediate_distill._lookup_layer(
    hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    idx: int,
    side: str
) -> torch.Tensor
```

```python
nemo_automodel.components.loss.intermediate_distill._masked_per_token_loss(
    s: torch.Tensor,
    t: torch.Tensor,
    attention_mask: torch.Tensor | None,
    loss_type: str
) -> torch.Tensor
```

```python
nemo_automodel.components.loss.intermediate_distill._pick_projector(
    projector: nemo_automodel.components.loss.intermediate_distill.ProjectorLike | None,
    s_idx: int
) -> collections.abc.Callable[[torch.Tensor], torch.Tensor] | None
```

```python
nemo_automodel.components.loss.intermediate_distill.intermediate_loss_function(
    student_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    teacher_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    layer_pairs: collections.abc.Sequence[tuple[int, int]],
    attention_mask: torch.Tensor | None = None,
    projector: nemo_automodel.components.loss.intermediate_distill.ProjectorLike | None = None,
    loss_type: str = 'mse',
    normalize: bool = False,
    layer_weights: collections.abc.Sequence[float] | None = None,
    reduction: str = 'mean'
) -> torch.Tensor
```

Per-token feature distillation between matched student/teacher layers.

```python
nemo_automodel.components.loss.intermediate_distill.intermediate_loss_pair(
    s_q_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    t_q_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    s_d_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    t_d_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    layer_pairs: collections.abc.Sequence[tuple[int, int]],
    attn_q: torch.Tensor | None = None,
    attn_d: torch.Tensor | None = None,
    kwargs = {}
) -> torch.Tensor
```

Sum of intermediate loss on query and doc sides.

```python
nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike = Sequence[torch.Tensor] | Mapping[int, torch.Tensor]
```

```python
nemo_automodel.components.loss.intermediate_distill.ProjectorLike = Callable[[torch.Tensor], torch.Tensor] | Mapping[int, Callable[[torch.Tensor], t...
```