class_docstr: The update block which contains the motion encoder and grus
  
  It must expose a ``hidden_dims`` attribute which is the hidden dimension size of its gru blocks
class_docstr_tok: ["The","update","block","which","contains","the","motion","encoder","and","grus","It","must","expose","a","``","hidden_dims","``","attribute","which","is","the","hidden","dimension","size","of","its","gru","blocks"]
code: def forward(
      self,
      hidden_states: List[Tensor],
      contexts: List[List[Tensor]],
      corr_features: Tensor,
      disparity: Tensor,
      level_processed: List[bool],
  ) -> List[Tensor]:
      # We call it reverse_i because it has a reversed ordering compared to hidden_states
      # see self.grus on the constructor for more detail
      for reverse_i, gru in enumerate(self.grus):
          i = len(self.grus) - 1 - reverse_i
          if level_processed[i]:
              # X is concatenation of 2x downsampled hidden_dim (or motion_features if no bigger dim) with
              # upsampled hidden_dim (or nothing if not exist).
              if i == 0:
                  features = self.motion_encoder(
                      disparity, corr_features
                  )
              else:
                  # 2x downsampled features from larger hidden states
                  features = F.avg_pool2d(
                      hidden_states[i - 1],
                      kernel_size=3,
                      stride=2,
                      padding=1,
                  )
  
              if i < len(self.grus) - 1:
                  # Concat with 2x upsampled features from smaller hidden states
                  _, _, h, w = hidden_states[i + 1].shape
                  features = torch.cat(
                      [
                          features,
                          F.interpolate(
                              hidden_states[i + 1],
                              size=(2 * h, 2 * w),
                              mode="bilinear",
                              align_corners=True,
                          ),
                      ],
                      dim=1,
                  )
  
              hidden_states[i] = gru(
                  hidden_states[i], features, contexts[i]
              )
  
              # NOTE: For slow-fast gru, we don't always want to calculate delta disparity for every call on UpdateBlock
              # Hence we move the delta disparity calculation to the RAFT-Stereo main forward
  
      return hidden_states
  
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 MultiLevelUpdateBlock obj's underlying member method code with example usage: obj.forward(hidden_states, contexts, corr_features, disparity, level_processed) and returns: hidden_states