def forward(
        self,
        features: dict[FeatureType, torch.Tensor] | torch.Tensor,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None:
        """Forward pass that either updates or predicts based on training status.

        Args:
            features (dict[FeatureType, torch.Tensor] | torch.Tensor): Input
                features. Can be either a dictionary mapping feature types to
                tensors, or a single tensor.

        Returns:
            torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None: During
                training, returns ``None`` after updating. During inference,
                returns density predictions.
        """
        if self.training:
            self.update(features)
            return None
        return self.predict(features)
Forward pass that either updates or predicts based on training status.

        Args:
            features (dict[FeatureType, torch.Tensor] | torch.Tensor): Input
                features. Can be either a dictionary mapping feature types to
                tensors, or a single tensor.

        Returns:
            torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None: During
                training, returns ``None`` after updating. During inference,
                returns density predictions.
        
src/anomalib/models/video/ai_vad/density.py