def forward(
        self,
        teacher_features: dict[str, torch.Tensor],
        student_features: dict[str, torch.Tensor],
    ) -> torch.Tensor:
        """Compute total loss across all feature layers.

        The total loss is computed as the sum of individual layer losses. Each
        layer loss measures the discrepancy between teacher and student features
        at that layer.

        Args:
            teacher_features (dict[str, torch.Tensor]): Dictionary mapping layer
                names to teacher feature tensors
            student_features (dict[str, torch.Tensor]): Dictionary mapping layer
                names to student feature tensors

        Returns:
            torch.Tensor: Total loss summed across all layers
        """
        layer_losses: list[torch.Tensor] = []
        for layer in teacher_features:
            loss = self.compute_layer_loss(teacher_features[layer], student_features[layer])
            layer_losses.append(loss)

        return torch.stack(layer_losses).sum()
Compute total loss across all feature layers.

        The total loss is computed as the sum of individual layer losses. Each
        layer loss measures the discrepancy between teacher and student features
        at that layer.

        Args:
            teacher_features (dict[str, torch.Tensor]): Dictionary mapping layer
                names to teacher feature tensors
            student_features (dict[str, torch.Tensor]): Dictionary mapping layer
                names to student feature tensors

        Returns:
            torch.Tensor: Total loss summed across all layers
        
src/anomalib/models/image/stfpm/loss.py