class_docstr: The motion encoder, part of the update block.
  
  Takes the current predicted flow and the correlation features as input and returns an encoded version of these.
class_docstr_tok: ["The","motion","encoder",",","part","of","the","update","block",".","Takes","the","current","predicted","flow","and","the","correlation","features","as","input","and","returns","an","encoded","version","of","these","."]
code: def __init__(
      self,
      *,
      in_channels_corr,
      corr_layers=(256, 192),
      flow_layers=(128, 64),
      out_channels=128,
  ):
      super().__init__()
  
      if len(flow_layers) != 2:
          raise ValueError(
              f"The expected number of flow_layers is 2, instead got {len(flow_layers)}"
          )
      if len(corr_layers) not in (1, 2):
          raise ValueError(
              f"The number of corr_layers should be 1 or 2, instead got {len(corr_layers)}"
          )
  
      self.convcorr1 = Conv2dNormActivation(
          in_channels_corr,
          corr_layers[0],
          norm_layer=None,
          kernel_size=1,
      )
      if len(corr_layers) == 2:
          self.convcorr2 = Conv2dNormActivation(
              corr_layers[0],
              corr_layers[1],
              norm_layer=None,
              kernel_size=3,
          )
      else:
          self.convcorr2 = nn.Identity()
  
      self.convflow1 = Conv2dNormActivation(
          2, flow_layers[0], norm_layer=None, kernel_size=7
      )
      self.convflow2 = Conv2dNormActivation(
          flow_layers[0], flow_layers[1], norm_layer=None, kernel_size=3
      )
  
      # out_channels - 2 because we cat the flow (2 channels) at the end
      self.conv = Conv2dNormActivation(
          corr_layers[-1] + flow_layers[-1],
          out_channels - 2,
          norm_layer=None,
          kernel_size=3,
      )
  
      self.out_channels = out_channels
  
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 1 to create a new MotionEncoder object from inherited base classes: nn with example: obj = MotionEncoder()