Source code for mednet.models.classify.training_steps

# SPDX-FileCopyrightText: Copyright © 2026 Idiap Research Institute <contact@idiap.ch>
#
# SPDX-License-Identifier: GPL-3.0-or-later

from ..model import Model




def default_training_step(model: Model, batch, batch_idx):
    """Perform a training step.

    Parameters
    ----------
    model
        Instance of `mednet.models.classify.model.Model` to use the training step with.
    batch
        The batch to apply the training step on. Should contain `image` and `target` keys.
    batch_idx
        The index of the batch, will be ignored.

    Returns
    -------
        The value of training loss for the batch.
    """
    del batch_idx
    return model.train_loss(
        model(model.augmentation_transforms(batch["image"])),
        batch["target"],
    )





def fairness_aware_training_step(model: Model, batch, batch_idx):
    """Perform a fairnesss-aware training step, for use with losses expecting a target and sensitive attributes.

    Parameters
    ----------
    model
        Instance of `mednet.models.classify.model.Model` to use the training step with.
    batch
        The batch to apply the training step on. Should contain `image`, `target` and `sensitive_attr` keys.
    batch_idx
        The index of the batch, will be ignored.

    Returns
    -------
        The value of training loss for the batch.
    """
    del batch_idx
    return model.train_loss(
        model(model.augmentation_transforms(batch["image"])),
        batch["target"],
        batch["sensitive_attr"],
    )