def forward(self, first_frame: torch.Tensor, last_frame: torch.Tensor) -> torch.Tensor:
        """Forward pass through the flow extractor.

        Args:
            first_frame (torch.Tensor): Batch of starting frames of shape (N, 3, H, W).
            last_frame (torch.Tensor): Batch of last frames of shape (N, 3, H, W).

        Returns:
            Tensor: Estimated optical flow map of shape (N, 2, H, W).
        """
        height, width = first_frame.shape[-2:]

        # preprocess batch
        first_frame, last_frame = self.pre_process(first_frame, last_frame)

        # get flow maps
        with torch.no_grad():
            flows = self.model(first_frame, last_frame)[-1]

        # convert back to original size
        return F.resize(flows, [height, width], antialias=False)
Forward pass through the flow extractor.

        Args:
            first_frame (torch.Tensor): Batch of starting frames of shape (N, 3, H, W).
            last_frame (torch.Tensor): Batch of last frames of shape (N, 3, H, W).

        Returns:
            Tensor: Estimated optical flow map of shape (N, 2, H, W).
        
src/anomalib/models/video/ai_vad/flow.py