def forward(
        self,
        batch: torch.Tensor,
        anomaly_map_to_generate: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor] | InferenceBatch:
        """Forward pass through the model.

        Args:
            batch (torch.Tensor): Input batch of images.
            anomaly_map_to_generate (torch.Tensor | None, optional): Anomaly map
                to use for generating quantized defects. Should be ``None`` if
                not in training phase 2. Defaults to ``None``.

        Returns:
            dict[str, torch.Tensor] | InferenceBatch: Output depends on mode:

            If testing:
                - ``anomaly_map``: Upsampled anomaly map
                - ``pred_score``: Image anomaly score

            If training phase 2:
                - ``recon_feat_hi``: Reconstructed non-quantized hi features
                  (F~_hi)
                - ``recon_feat_lo``: Reconstructed non-quantized lo features
                  (F~_lo)
                - ``embedding_bot``: Quantized features of non defective img
                  (Q_hi)
                - ``embedding_top``: Quantized features of non defective img
                  (Q_lo)
                - ``obj_spec_image``: Object-specific-decoded image (I_spc)
                - ``anomaly_map``: Predicted segmentation mask (M)
                - ``true_mask``: Resized ground-truth anomaly map (M_gt)

            If training phase 3:
                - ``anomaly_map``: Reconstructed anomaly map

        Raises:
            RuntimeError: If ``anomaly_map_to_generate`` is provided when not in
                training mode.

        Example:
            >>> model = DsrModel()
            >>> input_tensor = torch.randn(32, 3, 256, 256)
            >>> output = model(input_tensor)
            >>> output["anomaly_map"].shape
            torch.Size([32, 256, 256])
        """
        # Generate latent embeddings decoded image via general object decoder
        if anomaly_map_to_generate is None:
            # either evaluating or training phase 3
            with torch.no_grad():
                latent_model_outputs = self.discrete_latent_model(batch)
                gen_image = latent_model_outputs["recon_image"]
                embd_top = latent_model_outputs["quantized_t"]
                embd_bot = latent_model_outputs["quantized_b"]

                # Get embedders from the discrete latent model
                embedder_bot = self.discrete_latent_model.vq_vae_bot
                embedder_top = self.discrete_latent_model.vq_vae_top

                # Copy embeddings in order to input them to the subspace
                # restriction module
                anomaly_embedding_bot_copy = embd_bot.clone()
                anomaly_embedding_top_copy = embd_top.clone()

                # Apply subspace restriction module to copied embeddings
                _, recon_embd_bot = self.subspace_restriction_module_hi(anomaly_embedding_bot_copy, embedder_bot)
                _, recon_embd_top = self.subspace_restriction_module_lo(anomaly_embedding_top_copy, embedder_top)

                # Upscale top (lo) embedding
                up_quantized_recon_t = self.discrete_latent_model.upsample_t(recon_embd_top)

                # Concat embeddings and reconstruct image (object specific
                # decoder)
                quant_join = torch.cat((up_quantized_recon_t, recon_embd_bot), dim=1)
                obj_spec_image = self.image_reconstruction_network(quant_join)

                # Anomaly detection module
                out_mask = self.anomaly_detection_module(obj_spec_image, gen_image)
                out_mask_sm = torch.softmax(out_mask, dim=1)

            # Mask upsampling and score calculation
            upsampled_mask = self.upsampling_module(obj_spec_image, gen_image, out_mask_sm)
            out_mask_sm_up = torch.softmax(upsampled_mask, dim=1)

            # if training phase 3, return upsampled softmax mask
            if self.training:
                return {"anomaly_map": out_mask_sm_up}
            # if testing, extract image score
            out_mask_averaged = torch.nn.functional.avg_pool2d(
                out_mask_sm[:, 1:, :, :],
                21,
                stride=1,
                padding=21 // 2,
            ).detach()
            image_score = torch.amax(out_mask_averaged, dim=(2, 3)).squeeze()

            # prevent crash when image_score is a single value (batch size of 1)
            if image_score.size() == torch.Size([]):
                image_score = image_score.unsqueeze(0)

            anomaly_map = out_mask_sm_up[:, 1, :, :]

            return InferenceBatch(pred_score=image_score, anomaly_map=anomaly_map)

        if anomaly_map_to_generate is not None and self.training:
            # we are in phase two

            # Generate anomaly strength factors
            anom_str_lo = (
                torch.rand(batch.shape[0]) * (1.0 - self.latent_anomaly_strength) + self.latent_anomaly_strength
            ).cuda()
            anom_str_hi = (
                torch.rand(batch.shape[0]) * (1.0 - self.latent_anomaly_strength) + self.latent_anomaly_strength
            ).cuda()

            # Generate image through general object decoder, and defective & non
            # defective quantized feature maps.
            with torch.no_grad():
                latent_model_outputs = self.discrete_latent_model(
                    batch,
                    anomaly_map_to_generate,
                    anom_str_lo,
                    anom_str_hi,
                )
            gen_image_def = latent_model_outputs["recon_image"]
            true_anomaly_map = latent_model_outputs["anomaly_mask"]
            embd_top = latent_model_outputs["quantized_t"]
            embd_bot = latent_model_outputs["quantized_b"]
            embd_top_def = latent_model_outputs["anomaly_embedding_lo"]
            embd_bot_def = latent_model_outputs["anomaly_embedding_hi"]

            # Restore the features to normality with the Subspace restriction
            # modules
            recon_feat_hi, recon_embeddings_hi = self.subspace_restriction_module_hi(
                embd_bot_def,
                self.discrete_latent_model.vq_vae_bot,
            )
            recon_feat_lo, recon_embeddings_lo = self.subspace_restriction_module_lo(
                embd_top_def,
                self.discrete_latent_model.vq_vae_top,
            )

            # Reconstruct the image from the reconstructed features
            # with the object-specific image reconstruction module
            up_quantized_recon_t = self.discrete_latent_model.upsample_t(recon_embeddings_lo)
            quant_join = torch.cat((up_quantized_recon_t, recon_embeddings_hi), dim=1)
            spec_image_def = self.image_reconstruction_network(quant_join)

            # Generate the anomaly segmentation map
            out_mask = self.anomaly_detection_module(spec_image_def.detach(), gen_image_def.detach())
            out_mask_sm = torch.softmax(out_mask, dim=1)

            # Outputs
            return {
                "recon_feat_hi": recon_feat_hi,
                "recon_feat_lo": recon_feat_lo,
                "embedding_bot": embd_bot,
                "embedding_top": embd_top,
                "obj_spec_image": spec_image_def,
                "anomaly_map": out_mask_sm,
                "true_anomaly_map": true_anomaly_map,
            }
        msg = "There should not be an anomaly map to generate when not training"
        raise RuntimeError(msg)
Forward pass through the model.

        Args:
            batch (torch.Tensor): Input batch of images.
            anomaly_map_to_generate (torch.Tensor | None, optional): Anomaly map
                to use for generating quantized defects. Should be ``None`` if
                not in training phase 2. Defaults to ``None``.

        Returns:
            dict[str, torch.Tensor] | InferenceBatch: Output depends on mode:

            If testing:
                - ``anomaly_map``: Upsampled anomaly map
                - ``pred_score``: Image anomaly score

            If training phase 2:
                - ``recon_feat_hi``: Reconstructed non-quantized hi features
                  (F~_hi)
                - ``recon_feat_lo``: Reconstructed non-quantized lo features
                  (F~_lo)
                - ``embedding_bot``: Quantized features of non defective img
                  (Q_hi)
                - ``embedding_top``: Quantized features of non defective img
                  (Q_lo)
                - ``obj_spec_image``: Object-specific-decoded image (I_spc)
                - ``anomaly_map``: Predicted segmentation mask (M)
                - ``true_mask``: Resized ground-truth anomaly map (M_gt)

            If training phase 3:
                - ``anomaly_map``: Reconstructed anomaly map

        Raises:
            RuntimeError: If ``anomaly_map_to_generate`` is provided when not in
                training mode.

        Example:
            >>> model = DsrModel()
            >>> input_tensor = torch.randn(32, 3, 256, 256)
            >>> output = model(input_tensor)
            >>> output["anomaly_map"].shape
            torch.Size([32, 256, 256])
        
src/anomalib/models/image/dsr/torch_model.py