class_docstr: The update block which contains the motion encoder, the recurrent block, and the flow head.
  
  It must expose a ``hidden_state_size`` attribute which is the hidden state size of its recurrent block.
class_docstr_tok: ["The","update","block","which","contains","the","motion","encoder",",","the","recurrent","block",",","and","the","flow","head",".","It","must","expose","a","``","hidden_state_size","``","attribute","which","is","the","hidden","state","size","of","its","recurrent","block","."]
code: def forward(self, hidden_state, context, corr_features, flow):
      motion_features = self.motion_encoder(flow, corr_features)
      x = torch.cat([context, motion_features], dim=1)
  
      hidden_state = self.recurrent_block(hidden_state, x)
      delta_flow = self.flow_head(hidden_state)
      return hidden_state, delta_flow
  
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 UpdateBlock obj's underlying member method code with example usage: obj.forward(hidden_state, context, corr_features, flow) and returns: hidden_state,  delta_flow