def training_step(self, batch: Batch, *args, **kwargs) -> STEP_OUTPUT:
        """Perform training step for DRAEM.

        The step consists of:
        1. Generating simulated anomalies
        2. Computing reconstructions and predictions
        3. Calculating the loss

        Args:
            batch (Batch): Input batch containing images and metadata.
            args: Additional positional arguments (unused).
            kwargs: Additional keyword arguments (unused).

        Returns:
            STEP_OUTPUT: Dictionary containing the training loss.
        """
        del args, kwargs  # These variables are not used.

        input_image = batch.image
        # Apply corruption to input image
        augmented_image, anomaly_mask = self.augmenter(input_image)
        # Generate model prediction
        reconstruction, prediction = self.model(augmented_image)
        # Compute loss
        loss = self.loss(input_image, reconstruction, anomaly_mask, prediction)

        if self.sspcab:
            loss += self.sspcab_lambda * self.sspcab_loss(
                self.sspcab_activations["input"],
                self.sspcab_activations["output"],
            )

        self.log("train_loss", loss.item(), on_epoch=True, prog_bar=True, logger=True)
        return {"loss": loss}
Perform training step for DRAEM.

        The step consists of:
        1. Generating simulated anomalies
        2. Computing reconstructions and predictions
        3. Calculating the loss

        Args:
            batch (Batch): Input batch containing images and metadata.
            args: Additional positional arguments (unused).
            kwargs: Additional keyword arguments (unused).

        Returns:
            STEP_OUTPUT: Dictionary containing the training loss.
        
src/anomalib/models/image/draem/lightning_model.py