def training_step(self, batch: Batch) -> STEP_OUTPUT:
        """Training Step of DSR.

        During the first phase, feeds the original image and simulated anomaly
        mask. During second phase, feeds a generated anomalous image to train
        the upsampling module.

        Args:
            batch (Batch): Input batch containing image, label and mask

        Returns:
            STEP_OUTPUT: Dictionary containing the loss value

        Example:
            >>> from anomalib.data import Batch
            >>> model = Dsr()
            >>> batch = Batch(
            ...     image=torch.randn(8, 3, 256, 256),
            ...     label=torch.zeros(8)
            ... )
            >>> output = model.training_step(batch)
            >>> isinstance(output, dict)
            True
            >>> "loss" in output
            True
        """
        ph1_opt, ph2_opt = self.optimizers()

        if self.current_epoch < self.second_phase:
            # we are not yet training the upsampling module: we are only using
            # the first optimizer
            input_image = batch.image
            # Create anomaly masks
            anomaly_mask = self.quantized_anomaly_generator.augment_batch(input_image)
            # Generate model prediction
            model_outputs = self.model(input_image, anomaly_mask)
            # Compute loss
            loss = self.second_stage_loss(
                model_outputs["recon_feat_hi"],
                model_outputs["recon_feat_lo"],
                model_outputs["embedding_bot"],
                model_outputs["embedding_top"],
                input_image,
                model_outputs["obj_spec_image"],
                model_outputs["anomaly_map"],
                model_outputs["true_anomaly_map"],
            )

            # compute manual optimizer step
            ph1_opt.zero_grad()
            self.manual_backward(loss)
            ph1_opt.step()

        else:
            # we are training the upsampling module
            input_image = batch.image
            # Generate anomalies
            input_image, anomaly_maps = self.perlin_generator(input_image)
            # Get model prediction
            model_outputs = self.model(input_image)
            # Calculate loss
            loss = self.third_stage_loss(model_outputs["anomaly_map"], anomaly_maps)

            # compute manual optimizer step
            ph2_opt.zero_grad()
            self.manual_backward(loss)
            ph2_opt.step()

        self.log("train_loss", loss.item(), on_epoch=True, prog_bar=True, logger=True)
        return {"loss": loss}
Training Step of DSR.

        During the first phase, feeds the original image and simulated anomaly
        mask. During second phase, feeds a generated anomalous image to train
        the upsampling module.

        Args:
            batch (Batch): Input batch containing image, label and mask

        Returns:
            STEP_OUTPUT: Dictionary containing the loss value

        Example:
            >>> from anomalib.data import Batch
            >>> model = Dsr()
            >>> batch = Batch(
            ...     image=torch.randn(8, 3, 256, 256),
            ...     label=torch.zeros(8)
            ... )
            >>> output = model.training_step(batch)
            >>> isinstance(output, dict)
            True
            >>> "loss" in output
            True
        
src/anomalib/models/image/dsr/lightning_model.py