def training_step(self, batch: Batch, *args, **kwargs) -> STEP_OUTPUT:  # noqa: ARG002 | unused arguments
        """Perform a training step.

        Args:
            batch (Batch): Input batch containing images
            *args: Variable length argument list
            **kwargs: Arbitrary keyword arguments

        Returns:
            STEP_OUTPUT: Dictionary containing the loss value
        """
        z, ljd = self.model(batch.image)
        loss = self.loss(z, ljd)
        self.log_dict({"loss": loss}, on_step=True, on_epoch=False, prog_bar=False, logger=True)
        return {"loss": loss}
Perform a training step.

        Args:
            batch (Batch): Input batch containing images
            *args: Variable length argument list
            **kwargs: Arbitrary keyword arguments

        Returns:
            STEP_OUTPUT: Dictionary containing the loss value
        
src/anomalib/models/image/uflow/lightning_model.py