Class ESPnetLightningModule

espnet3/components/modeling/lightning_module.py:28–1124

ESPnet3 LightningModule wrapper for model training and data integration. This wrapper keeps the common ESPnet3 model contract unchanged: ```python loss, stats, weight = model(**batch) ``` Most models should continue to return a single scalar loss tensor. The training loop

class ESPnetLightningModule(lightning.LightningModule):
    """ESPnet3 LightningModule wrapper for model training and data integration.

    This wrapper keeps the common ESPnet3 model contract unchanged:

    ```python
    loss, stats, weight = model(**batch)
    ```

    Most models should continue to return a single scalar loss tensor. The training
    loop then behaves exactly like conventional Lightning single-optimizer training.

    When multiple optimizers are configured, the same return value is expected,
    but the `loss` field must carry optimizer routing information through
    `OptimizationStep`.
    The model still returns one `stats` dict and one optional `weight` value; only
    the type of `loss` changes.

    Example:
        Single optimizer path:
        ```python
        def forward(self, **batch):
            loss = ...
            stats = {"loss": loss.detach(), "acc": acc.detach()}
            weight = torch.tensor(batch_size, device=loss.device)
            return loss, stats, weight
        ```

        GAN-style path updating both optimizers in a single batch:
        ```python
        def forward(self, **batch):
            g_loss = ...
            d_loss = ...
            stats = {
                "generator_loss": g_loss.detach(),
                "discriminator_loss": d_loss.detach(),
            }
            return [
                OptimizationStep(loss=g_loss, name="generator"),
                OptimizationStep(loss=d_loss, name="discriminator"),
            ], stats, None
        ```

        GAN-style path updating only the generator for one batch:
        ```python
        def forward(self, **batch):
            g_loss = ...
            stats = {"generator_loss": g_loss.detach()}
            return OptimizationStep(loss=g_loss, name="generator"), stats, None
        ```

    Notes:
        - Returning `OptimizationStep` with the single optimizer is forbidden.
        - In the multi-optimizer path, only optimizers named by returned
          `OptimizationStep` objects are touched for that batch. Optimizers omitted
          from the list are left untouched entirely.
        - The order of `OptimizationStep` entries is the exact backward/step order.
        - NaN or Inf in any returned loss causes the whole batch to be skipped on