def forward(
        self,
        batch: torch.Tensor,
        anomaly_mask: torch.Tensor | None = None,
        anom_str_lo: torch.Tensor | None = None,
        anom_str_hi: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor]:
        """Generate quantized feature maps.

        Generates quantized feature maps of batch of input images as well as their
        reconstruction based on the general appearance decoder.

        Args:
            batch (Tensor): Batch of input images.
            anomaly_mask (Tensor | None): Anomaly mask to be used to generate anomalies on
                the quantized feature maps.
            anom_str_lo (torch.Tensor | None): Strength of generated anomaly lo.
            anom_str_hi (torch.Tensor | None): Strength of generated anomaly hi.

        Returns:
            dict[str, torch.Tensor]:
            If generating an anomaly mask:
                - General object decoder-decoded anomalous image
                - Reshaped ground truth anomaly map
                - Non defective quantized lo feature
                - Non defective quantized hi feature
                - Non quantized subspace encoded defective lo feature
                - Non quantized subspace encoded defective hi feature
            Else:
                - General object decoder-decoded image
                - Quantized lo feature
                - Quantized hi feature
        """
        # Encoder Hi
        enc_b = self._encoder_b(batch)

        # Encoder Lo -- F_Lo
        enc_t = self._encoder_t(enc_b)
        zt = self._pre_vq_conv_top(enc_t)

        # Quantize F_Lo with K_Lo
        quantized_t = self._vq_vae_top(zt)

        # Upsample Q_Lo
        up_quantized_t = self.upsample_t(quantized_t)

        # Concatenate and transform the output of Encoder_Hi and upsampled Q_lo -- F_Hi
        feat = torch.cat((enc_b, up_quantized_t), dim=1)
        zb = self._pre_vq_conv_bot(feat)

        # Quantize F_Hi with K_Hi
        quantized_b = self._vq_vae_bot(zb)

        # generate anomalies
        anomaly_embedding_hi = None
        anomaly_embedding_lo = None

        # define outputs
        outputs = {"quantized_b": quantized_b, "quantized_t": quantized_t}

        if anomaly_mask is not None:
            # Generate feature-based anomalies on F_lo
            anomaly_embedding_lo = self.generate_fake_anomalies_joined(
                zt,
                quantized_t,
                self._vq_vae_top.embedding.weight,
                anomaly_mask,
                anom_str_lo,
            )

            up_quantized_t_defect = self.upsample_t(anomaly_embedding_lo)
            feat_defect = torch.cat((enc_b, up_quantized_t_defect), dim=1)
            zb_defect = self._pre_vq_conv_bot(feat_defect)
            quantized_b_defect = self._vq_vae_bot(zb_defect)

            # Generate feature-based anomalies on F_hi
            anomaly_embedding_hi = self.generate_fake_anomalies_joined(
                zb_defect,
                quantized_b_defect,
                self._vq_vae_bot.embedding.weight,
                anomaly_mask,
                anom_str_hi,
            )

            # get anomaly embeddings
            use_both = torch.randint(0, 2, (batch.shape[0], 1, 1, 1)).cuda().float()
            use_lo = torch.randint(0, 2, (batch.shape[0], 1, 1, 1)).cuda().float()
            use_hi = 1 - use_lo

            anomaly_embedding_hi_usebot = self.generate_fake_anomalies_joined(
                zb,
                quantized_b,
                self._vq_vae_bot.embedding.weight,
                anomaly_mask,
                anom_str_hi,
            )

            anomaly_embedding_lo_usebot = quantized_t
            anomaly_embedding_hi_usetop = quantized_b
            anomaly_embedding_lo_usetop = anomaly_embedding_lo
            anomaly_embedding_hi_not_both = use_hi * anomaly_embedding_hi_usebot + use_lo * anomaly_embedding_hi_usetop
            anomaly_embedding_lo_not_both = use_hi * anomaly_embedding_lo_usebot + use_lo * anomaly_embedding_lo_usetop
            anomaly_embedding_hi = (
                (anomaly_embedding_hi * use_both + anomaly_embedding_hi_not_both * (1.0 - use_both)).detach().clone()
            )
            anomaly_embedding_lo = (
                (anomaly_embedding_lo * use_both + anomaly_embedding_lo_not_both * (1.0 - use_both)).detach().clone()
            )

            anomaly_embedding_hi_copy = anomaly_embedding_hi.clone()
            anomaly_embedding_lo_copy = anomaly_embedding_lo.clone()

            # apply the general appearance decoder to the anomaly embeddings
            up_quantized_anomaly_t = self.upsample_t(anomaly_embedding_lo_copy)
            quant_join_anomaly = torch.cat((up_quantized_anomaly_t, anomaly_embedding_hi_copy), dim=1)
            recon_image = self._decoder_b(quant_join_anomaly)

            # Resize the ground truth anomaly map to closely match the augmented features
            down_ratio_x_hi = int(anomaly_mask.shape[3] / quantized_b_defect.shape[3])
            anomaly_mask_hi = torch.nn.functional.max_pool2d(anomaly_mask, (down_ratio_x_hi, down_ratio_x_hi)).float()
            anomaly_mask_hi = torch.nn.functional.interpolate(anomaly_mask_hi, scale_factor=down_ratio_x_hi)
            down_ratio_x_lo = int(anomaly_mask.shape[3] / quantized_t.shape[3])
            anomaly_mask_lo = torch.nn.functional.max_pool2d(anomaly_mask, (down_ratio_x_lo, down_ratio_x_lo)).float()
            anomaly_mask_lo = torch.nn.functional.interpolate(anomaly_mask_lo, scale_factor=down_ratio_x_lo)
            anomaly_mask = anomaly_mask_lo * use_both + (anomaly_mask_lo * use_lo + anomaly_mask_hi * use_hi) * (
                1.0 - use_both
            )

            # reminder : top = lo, bot = hi!
            outputs["recon_image"] = recon_image
            outputs["anomaly_mask"] = anomaly_mask
            outputs["anomaly_embedding_lo"] = anomaly_embedding_lo
            outputs["anomaly_embedding_hi"] = anomaly_embedding_hi

        else:
            # Concatenate Q_Hi and Q_Lo and input it into the General appearance decoder
            quant_join = torch.cat((up_quantized_t, quantized_b), dim=1)
            recon_image = self._decoder_b(quant_join)

            outputs["recon_image"] = recon_image

        return outputs
Generate quantized feature maps.

        Generates quantized feature maps of batch of input images as well as their
        reconstruction based on the general appearance decoder.

        Args:
            batch (Tensor): Batch of input images.
            anomaly_mask (Tensor | None): Anomaly mask to be used to generate anomalies on
                the quantized feature maps.
            anom_str_lo (torch.Tensor | None): Strength of generated anomaly lo.
            anom_str_hi (torch.Tensor | None): Strength of generated anomaly hi.

        Returns:
            dict[str, torch.Tensor]:
            If generating an anomaly mask:
                - General object decoder-decoded anomalous image
                - Reshaped ground truth anomaly map
                - Non defective quantized lo feature
                - Non defective quantized hi feature
                - Non quantized subspace encoded defective lo feature
                - Non quantized subspace encoded defective hi feature
            Else:
                - General object decoder-decoded image
                - Quantized lo feature
                - Quantized hi feature
        
src/anomalib/models/image/dsr/torch_model.py