def get_features(self, batch: torch.Tensor) -> torch.Tensor:
        """Extract features from the pretrained network.

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

        Returns:
            Union[torch.Tensor, Tuple[torch.Tensor, torch.Size]]: Features during
                training, or tuple of (features, feature_shapes) during inference.
        """
        self.feature_extractor.eval()
        features = self.feature_extractor(batch)[self.layer]
        batch_size = len(features)
        if self.pooling_kernel_size > 1:
            features = F.avg_pool2d(input=features, kernel_size=self.pooling_kernel_size)
        feature_shapes = features.shape
        features = features.view(batch_size, -1).detach()
        return features if self.training else (features, feature_shapes)
Extract features from the pretrained network.

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

        Returns:
            Union[torch.Tensor, Tuple[torch.Tensor, torch.Size]]: Features during
                training, or tuple of (features, feature_shapes) during inference.
        
src/anomalib/models/image/dfm/torch_model.py