docstr: Return disparity predictions on every iteration as a list of Tensor.
  args:
      left_image (Tensor): The input left image with layout B, C, H, W
      right_image (Tensor): The input right image with layout B, C, H, W
      flow_init (Optional[Tensor]): Initial estimate for the disparity. Default: None
      num_iters (int): Number of update block iteration on the largest resolution. Default: 12
code: def forward(
      self,
      left_image: Tensor,
      right_image: Tensor,
      flow_init: Optional[Tensor] = None,
      num_iters: int = 12,
  ) -> List[Tensor]:
      """
      Return disparity predictions on every iteration as a list of Tensor.
      args:
          left_image (Tensor): The input left image with layout B, C, H, W
          right_image (Tensor): The input right image with layout B, C, H, W
          flow_init (Optional[Tensor]): Initial estimate for the disparity. Default: None
          num_iters (int): Number of update block iteration on the largest resolution. Default: 12
      """
      batch_size, _, h, w = left_image.shape
      torch._assert(
          (h, w) == right_image.shape[-2:],
          f"input images should have the same shape, instead got ({h}, {w}) != {right_image.shape[-2:]}",
      )
  
      torch._assert(
          (
              h % self.base_downsampling_ratio == 0
              and w % self.base_downsampling_ratio == 0
          ),
          f"input image H and W should be divisible by {self.base_downsampling_ratio}, instead got H={h} and W={w}",
      )
  
      fmaps = self.feature_encoder(
          torch.cat([left_image, right_image], dim=0)
      )
      fmap1, fmap2 = torch.chunk(fmaps, chunks=2, dim=0)
      torch._assert(
          fmap1.shape[-2:]
          == (
              h // self.base_downsampling_ratio,
              w // self.base_downsampling_ratio,
          ),
          f"The feature encoder should downsample H and W by {self.base_downsampling_ratio}",
      )
  
      corr_pyramid = self.corr_pyramid(fmap1, fmap2)
  
      # Multi level contexts
      context_outs = self.context_encoder(left_image)
  
      hidden_dims = self.update_block.hidden_dims
      context_out_channels = [
          context_outs[i].shape[1] - hidden_dims[i]
          for i in range(len(context_outs))
      ]
      hidden_states: List[Tensor] = []
      contexts: List[List[Tensor]] = []
      for i, context_conv in enumerate(self.context_convs):
          # As in the original paper, the actual output of the context encoder is split in 2 parts:
          # - one part is used to initialize the hidden state of the recurent units of the update block
          # - the rest is the "actual" context.
          hidden_state, context = torch.split(
              context_outs[i],
              [hidden_dims[i], context_out_channels[i]],
              dim=1,
          )
          hidden_states.append(torch.tanh(hidden_state))
          contexts.append(
              # mypy is technically correct here. The return type of `torch.split` was incorrectly annotated with
              # `List[int]` although it should have been `Tuple[Tensor, ...]`. However, the latter is not supported by
              # JIT and thus we have to keep the wrong annotation here and silence mypy.
              torch.split(  # type: ignore[arg-type]
                  context_conv(F.relu(context)),
                  [hidden_dims[i], hidden_dims[i], hidden_dims[i]],
                  dim=1,
              )
          )
  
      _, Cf, Hf, Wf = fmap1.shape
      coords0 = make_coords_grid(batch_size, Hf, Wf).to(fmap1.device)
      coords1 = make_coords_grid(batch_size, Hf, Wf).to(fmap1.device)
  
      # We use flow_init for cascade inference
      if flow_init is not None:
          coords1 = coords1 + flow_init
  
      disparity_predictions = []
      for _ in range(num_iters):
          coords1 = (
              coords1.detach()
          )  # Don't backpropagate gradients through this branch, see paper
          corr_features = self.corr_block(
              centroids_coords=coords1, corr_pyramid=corr_pyramid
          )
  
          disparity = coords1 - coords0
          if self.slow_fast:
              # Using slow_fast GRU (see paper section 3.4). The lower resolution are processed more often
              for i in range(1, self.num_level):
                  # We only processed the smallest i levels
                  level_processed = [False] * (self.num_level - i) + [
                      True
                  ] * i
                  hidden_states = self.update_block(
                      hidden_states,
                      contexts,
                      corr_features,
                      disparity,
                      level_processed=level_processed,
                  )
          hidden_states = self.update_block(
              hidden_states,
              contexts,
              corr_features,
              disparity,
              level_processed=[True] * self.num_level,
          )
          # Take the largest hidden_state to get the disparity
          hidden_state = hidden_states[0]
          delta_disparity = self.disparity_head(hidden_state)
          # in stereo mode, project disparity onto epipolar
          delta_disparity[:, 1] = 0.0
  
          coords1 = coords1 + delta_disparity
          up_mask = (
              None
              if self.mask_predictor is None
              else self.mask_predictor(hidden_state)
          )
          upsampled_disparity = upsample_flow(
              (coords1 - coords0),
              up_mask=up_mask,
              factor=self.base_downsampling_ratio,
          )
          disparity_predictions.append(upsampled_disparity[:, :1])
  
      return disparity_predictions
  
imports: from functools import partial
  from typing import Callable, List, Optional, Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  import torchvision.models.optical_flow.raft
  from torch import Tensor
  from torchvision.models._api import register_model, Weights, WeightsEnum
  from torchvision.models._utils import handle_legacy_interface
  from torchvision.models.optical_flow._utils import grid_sample, make_coords_grid, upsample_flow
  from torchvision.models.optical_flow.raft import FlowHead, MotionEncoder, ResidualBlock
  from torchvision.ops import Conv2dNormActivation
  from torchvision.prototype.transforms._presets import StereoMatching
  from torchvision.utils import _log_api_usage_once
label_desc: Use image node_id 2 for calling the RaftStereo obj's underlying member method code with example usage: obj.forward(left_image, right_image, flow_init, num_iters) and returns: disparity_predictions