code: def forward(self, image1, image2, num_flow_updates: int = 12):
      batch_size, _, h, w = image1.shape
      if (h, w) != image2.shape[-2:]:
          raise ValueError(
              f"input images should have the same shape, instead got ({h}, {w}) != {image2.shape[-2:]}"
          )
      if not (h % 8 == 0) and (w % 8 == 0):
          raise ValueError(
              f"input image H and W should be divisible by 8, instead got {h} (h) and {w} (w)"
          )
  
      fmaps = self.feature_encoder(torch.cat([image1, image2], dim=0))
      fmap1, fmap2 = torch.chunk(fmaps, chunks=2, dim=0)
      if fmap1.shape[-2:] != (h // 8, w // 8):
          raise ValueError(
              "The feature encoder should downsample H and W by 8"
          )
  
      self.corr_block.build_pyramid(fmap1, fmap2)
  
      context_out = self.context_encoder(image1)
      if context_out.shape[-2:] != (h // 8, w // 8):
          raise ValueError(
              "The context encoder should downsample H and W by 8"
          )
  
      # 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_size = self.update_block.hidden_state_size
      out_channels_context = context_out.shape[1] - hidden_state_size
      if out_channels_context <= 0:
          raise ValueError(
              f"The context encoder outputs {context_out.shape[1]} channels, but it should have at strictly more than hidden_state={hidden_state_size} channels"
          )
      hidden_state, context = torch.split(
          context_out, [hidden_state_size, out_channels_context], dim=1
      )
      hidden_state = torch.tanh(hidden_state)
      context = F.relu(context)
  
      coords0 = make_coords_grid(batch_size, h // 8, w // 8).to(
          fmap1.device
      )
      coords1 = make_coords_grid(batch_size, h // 8, w // 8).to(
          fmap1.device
      )
  
      flow_predictions = []
      for _ in range(num_flow_updates):
          coords1 = (
              coords1.detach()
          )  # Don't backpropagate gradients through this branch, see paper
          corr_features = self.corr_block.index_pyramid(
              centroids_coords=coords1
          )
  
          flow = coords1 - coords0
          hidden_state, delta_flow = self.update_block(
              hidden_state, context, corr_features, flow
          )
  
          coords1 = coords1 + delta_flow
  
          up_mask = (
              None
              if self.mask_predictor is None
              else self.mask_predictor(hidden_state)
          )
          upsampled_flow = upsample_flow(
              flow=(coords1 - coords0), up_mask=up_mask
          )
          flow_predictions.append(upsampled_flow)
  
      return flow_predictions
  
imports: from typing import List, Optional
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from torch.nn.modules.batchnorm import BatchNorm2d
  from torch.nn.modules.instancenorm import InstanceNorm2d
  from torchvision.ops import Conv2dNormActivation
  from ...transforms._presets import OpticalFlow
  from ...utils import _log_api_usage_once
  from .._api import register_model, Weights, WeightsEnum
  from .._utils import handle_legacy_interface
  from ._utils import grid_sample, make_coords_grid, upsample_flow
label_desc: Use image node_id 2 for calling the RAFT obj's underlying member method code with example usage: obj.forward(image1, image2, num_flow_updates) and returns: flow_predictions