docstr: Args:
      source: tensor of shape (B, C, H, W)
      reference: tensor of shape (B, C, H, W)
      flow_pred: tensor of shape (B, 2, H, W)
      valid_mask: tensor of shape (B, H, W) or None
  
  Returns:
      photometric loss of shape
code: def forward(
      self,
      source: Tensor,
      reference: Tensor,
      flow_pred: Tensor,
      valid_mask: Optional[Tensor] = None,
  ):
      """
      Args:
          source: tensor of shape (B, C, H, W)
          reference: tensor of shape (B, C, H, W)
          flow_pred: tensor of shape (B, 2, H, W)
          valid_mask: tensor of shape (B, H, W) or None
  
      Returns:
          photometric loss of shape
  
      """
      torch._assert(
          source.ndim == 4,
          "FlowPhotoMetricLoss: source must have 4 dimensions, but got {}".format(
              source.ndim
          ),
      )
      torch._assert(
          reference.ndim == source.ndim,
          "FlowPhotoMetricLoss: source and other must have the same number of dimensions, but got {} and {}".format(
              source.ndim, reference.ndim
          ),
      )
      torch._assert(
          flow_pred.shape[1] == 2,
          "FlowPhotoMetricLoss: flow_pred must have 2 channels, but got {}".format(
              flow_pred.shape[1]
          ),
      )
      torch._assert(
          flow_pred.ndim == 4,
          "FlowPhotoMetricLoss: flow_pred must have 4 dimensions, but got {}".format(
              flow_pred.ndim
          ),
      )
  
      B, C, H, W = source.shape
      flow_channels = flow_pred.shape[1]
  
      max_displacements = []
      for dim in range(flow_channels):
          shape_index = -1 - dim
          max_displacements.append(
              int(
                  self._max_displacement_ratio
                  * source.shape[shape_index]
              )
          )
  
      # mask out all pixels that have larger flow than the max flow allowed
      max_flow_mask = torch.logical_and(
          *[
              flow_pred[:, dim, :, :] < max_displacements[dim]
              for dim in range(flow_channels)
          ]
      )
  
      if valid_mask is not None:
          valid_mask = torch.logical_and(
              valid_mask, max_flow_mask
          ).unsqueeze(1)
      else:
          valid_mask = max_flow_mask.unsqueeze(1)
  
      grid = make_coords_grid(B, H, W, device=str(source.device))
      resampled_grids = grid - flow_pred
      resampled_grids = resampled_grids.permute(0, 2, 3, 1)
      resampled_source = grid_sample(
          reference, resampled_grids, mode="bilinear"
      )
  
      # compute SSIM loss
      ssim_loss = self._ssim_loss(
          resampled_source * valid_mask, source * valid_mask
      )
      l1_loss = (
          (resampled_source * valid_mask - source * valid_mask)
          .abs()
          .mean(axis=(-3, -2, -1))
      )
      loss = self._L1_weight * l1_loss + self._SSIM_weight * ssim_loss
  
      return loss.mean()
  
imports: from typing import List, Optional
  import torch
  from torch import nn, Tensor
  from torch.nn import functional
  from torchvision.prototype.models.depth.stereo.raft_stereo import grid_sample, make_coords_grid
label_desc: Use image node_id 2 for calling the FlowPhotoMetricLoss obj's underlying member method code with example usage: obj.forward(source, reference, flow_pred, valid_mask) and returns: loss