def forward(
        self,
        images: torch.Tensor,
    ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]] | InferenceBatch:
        """Forward pass through teacher and student networks.

        The forward pass behavior differs between training and evaluation:
        - Training: Returns features from both teacher and student networks
        - Evaluation: Returns anomaly maps generated from feature differences

        Args:
            images (torch.Tensor): Batch of input images with shape
                ``(N, C, H, W)``.

        Returns:
            Training mode:
                tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
                    Features from teacher and student networks respectively.
                    Each dict maps layer names to feature tensors.
            Evaluation mode:
                InferenceBatch:
                    Batch containing anomaly maps and prediction scores.

        Example:
            >>> import torch
            >>> from anomalib.models.image.stfpm.torch_model import STFPMModel
            >>> model = STFPMModel(layers=["layer1", "layer2", "layer3"])
            >>> input_tensor = torch.randn(1, 3, 256, 256)
            >>> # Training mode
            >>> model.train()
            >>> teacher_feats, student_feats = model(input_tensor)
            >>> # Evaluation mode
            >>> model.eval()
            >>> predictions = model(input_tensor)
        """
        output_size = images.shape[-2:]
        if self.tiler:
            images = self.tiler.tile(images)
        teacher_features: dict[str, torch.Tensor] = self.teacher_model(images)
        student_features: dict[str, torch.Tensor] = self.student_model(images)

        if self.tiler:
            for layer, data in teacher_features.items():
                teacher_features[layer] = self.tiler.untile(data)
            for layer, data in student_features.items():
                student_features[layer] = self.tiler.untile(data)

        if self.training:
            return teacher_features, student_features

        anomaly_map = self.anomaly_map_generator(
            teacher_features=teacher_features,
            student_features=student_features,
            image_size=output_size,
        )
        pred_score = torch.amax(anomaly_map, dim=(-2, -1))
        return InferenceBatch(pred_score=pred_score, anomaly_map=anomaly_map)
Forward pass through teacher and student networks.

        The forward pass behavior differs between training and evaluation:
        - Training: Returns features from both teacher and student networks
        - Evaluation: Returns anomaly maps generated from feature differences

        Args:
            images (torch.Tensor): Batch of input images with shape
                ``(N, C, H, W)``.

        Returns:
            Training mode:
                tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
                    Features from teacher and student networks respectively.
                    Each dict maps layer names to feature tensors.
            Evaluation mode:
                InferenceBatch:
                    Batch containing anomaly maps and prediction scores.

        Example:
            >>> import torch
            >>> from anomalib.models.image.stfpm.torch_model import STFPMModel
            >>> model = STFPMModel(layers=["layer1", "layer2", "layer3"])
            >>> input_tensor = torch.randn(1, 3, 256, 256)
            >>> # Training mode
            >>> model.train()
            >>> teacher_feats, student_feats = model(input_tensor)
            >>> # Evaluation mode
            >>> model.eval()
            >>> predictions = model(input_tensor)
        
src/anomalib/models/image/stfpm/torch_model.py