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 __init__(
      self, *, motion_encoder: MotionEncoder, hidden_dims: List[int]
  ):
      super().__init__()
      self.motion_encoder = motion_encoder
  
      # The GRU input size is the size of previous level hidden_dim plus next level hidden_dim
      # if this is the first gru, then we replace previous level with motion_encoder output channels
      # for the last GRU, we don't add the next level hidden_dim
      gru_input_dims = []
      for i in range(len(hidden_dims)):
          input_dim = (
              hidden_dims[i - 1]
              if i > 0
              else motion_encoder.out_channels
          )
          if i < len(hidden_dims) - 1:
              input_dim += hidden_dims[i + 1]
          gru_input_dims.append(input_dim)
  
      self.grus = nn.ModuleList(
          [
              ConvGRU(
                  input_size=gru_input_dims[i],
                  hidden_size=hidden_dims[i],
                  kernel_size=3,
                  padding=1,
              )
              # Ideally we should reverse the direction during forward to use the gru with the smallest resolution
              # first however currently there is no way to reverse a ModuleList that is jit script compatible
              # hence we reverse the ordering of self.grus on the constructor instead
              # see: https://github.com/pytorch/pytorch/issues/31772
              for i in reversed(list(range(len(hidden_dims))))
          ]
      )
  
      self.hidden_dims = hidden_dims
  
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 1 to create a new MultiLevelUpdateBlock object from inherited base classes: nn with example: obj = MultiLevelUpdateBlock()