code: def compute_batch_loss(self, image1, image2):
      assert self.loss_functions.entropy_fn is not None
      assert image1.ndim == image2.ndim == 4
  
      flow, image2_est, probabilities = self.forward(image1, image2)
  
      # compute distortion loss
      distortion_loss = self.loss_functions.distortion_fn(
          image2, image2_est
      )
  
      # compute compression loss, average over num pixels
      num_pixels = image1.shape[0] * image1.shape[-2] * image1.shape[-1]
      entropy_loss = self.loss_functions.entropy_fn(
          probabilities, num_pixels
      )
  
      return (
          LossValues(
              distortion_loss=distortion_loss,
              flow_entropy_loss=entropy_loss,
              resid_entropy_loss=None,
          ),
          OutputTensors(
              flow=flow.detach(),
              image1=image1.detach(),
              image2=image2.detach(),
              image2_est=image2_est.detach(),
          ),
      )
  
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 3 for calling the DvcStage3 obj's underlying member method code with example usage: obj.compute_batch_loss(image1, image2) and returns: LossValues,  OutputTensors