code: def forward(self, image1, image2):
      flow = self.motion_estimator(image1, image2)
  
      latent, sizes = self.motion_encoder(flow)
      latent, probabilities = self.motion_entropy_bottleneck(latent)
      flow = self.motion_decoder(latent, sizes)
      image2_est = ncF.dense_image_warp(
          image1, flow.permute(0, 2, 3, 1)
      )
  
      return flow, image2_est, probabilities
  
imports: from typing import Callable, NamedTuple, Optional
  import torch
  import torch.nn
  from torch import Tensor
  import neuralcompression.functional
label_desc: Use image node_id 2 for calling the DvcStage2 obj's underlying member method code with example usage: obj.forward(image1, image2) and returns: flow,  image2_est,  probabilities