def forward(self, batch: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | InferenceBatch:
        """Forward pass through both sub-networks.

        Args:
            batch (torch.Tensor): Input batch of images of shape
                ``(batch_size, channels, height, width)``

        Returns:
            During training:
                tuple: Tuple containing:
                    - Reconstructed images
                    - Predicted anomaly masks
            During inference:
                InferenceBatch: Contains anomaly map and prediction score

        Example:
            >>> model = DraemModel()
            >>> batch = torch.randn(32, 3, 256, 256)
            >>> reconstruction, prediction = model(batch)  # Training mode
            >>> model.eval()
            >>> output = model(batch)  # Inference mode
            >>> assert isinstance(output, InferenceBatch)
        """
        reconstruction = self.reconstructive_subnetwork(batch)
        concatenated_inputs = torch.cat([batch, reconstruction], axis=1)
        prediction = self.discriminative_subnetwork(concatenated_inputs)
        if self.training:
            return reconstruction, prediction

        anomaly_map = torch.softmax(prediction, dim=1)[:, 1, ...]
        pred_score = torch.amax(anomaly_map, dim=(-2, -1))
        return InferenceBatch(pred_score=pred_score, anomaly_map=anomaly_map)
Forward pass through both sub-networks.

        Args:
            batch (torch.Tensor): Input batch of images of shape
                ``(batch_size, channels, height, width)``

        Returns:
            During training:
                tuple: Tuple containing:
                    - Reconstructed images
                    - Predicted anomaly masks
            During inference:
                InferenceBatch: Contains anomaly map and prediction score

        Example:
            >>> model = DraemModel()
            >>> batch = torch.randn(32, 3, 256, 256)
            >>> reconstruction, prediction = model(batch)  # Training mode
            >>> model.eval()
            >>> output = model(batch)  # Inference mode
            >>> assert isinstance(output, InferenceBatch)
        
src/anomalib/models/image/draem/torch_model.py