text: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from .hd3_ops import *
  
  
  class LossCalculator(object):
  
      def __init__(self, task):
          assert task in ['flow', 'stereo']
          self.task = task
          self.dim = 1 if task == 'stereo' else 2
  
      def __call__(self, outputs, inputs):
          ms_prob = outputs['ms_prob']
          ms_pred = outputs['ms_pred']
          corr_range = outputs['corr_range']
          ds = outputs['downsample']
          gt = inputs['flows'][:, 0]
          B, C, H, W = gt.size()
          lv = len(ms_prob)
          criterion = nn.KLDivLoss(reduction='batchmean').to(device=gt.device)
          losses = {}
          kld_loss = 0
          for l in range(lv):
              scaled_gt, valid_mask = downsample_flow(gt, 1 / 2**(ds - l))
              if self.task == 'stereo':
                  scaled_gt = scaled_gt[:, 0, :, :].unsqueeze(1)
              if l > 0:
                  scaled_gt = scaled_gt - F.interpolate(
                      ms_pred[l - 1],
                      scale_factor=2,
                      mode='bilinear',
                      align_corners=True)
              scaled_gt = scaled_gt / 2**(ds - l)
              gt_dist = vector2density(scaled_gt, corr_range[l],
                                       self.dim) * valid_mask
              kld_loss += 4**(ds - l) / (H * W) * criterion(
                  F.log_softmax(ms_prob[l], dim=1), gt_dist.detach())
  
          losses['loss'] = kld_loss
          for loss_type, loss_value in losses.items():
              losses[loss_type] = loss_value.reshape(1)
          return losses
  
  
  def EndPointError(output, gt):
      # output: [B, 1/2, H, W], stereo or flow prediction
      # gt: [B, C, H, W], 2D ground-truth annotation which may contain a mask
      # NOTE: To benchmark the result, please ensure the ground-truth keeps
      # its ORIGINAL RESOLUTION.
      if output.size(1) == 1:  # stereo
          output = disp2flow(output)
      output = resize_dense_vector(output, gt.size(2), gt.size(3))
      error = torch.norm(output - gt[:, :2, :, :], 2, 1, keepdim=False)
      if gt.size(1) == 3:
          mask = (gt[:, 2, :, :] > 0).float()
      else:
          mask = torch.ones_like(error)
      epe = (error * mask).sum() / mask.sum()
      return epe.reshape(1)
  