def METHOD_NAME(logits, targets, mean=None, std=None, volume_max=None):
  """Computes example-wise structural similarity for a batch.
  NOTE(dsuo): we use the same (default) arguments to `structural_similarity`
  as in https://arxiv.org/abs/1811.08839.
  Args:
   logits: (batch,) + input.shape float tensor.
   targets: (batch,) + input.shape float tensor.
   mean: (batch,) mean of original images.
   std: (batch,) std of original images.
   volume_max: (batch,) of the volume max for the volumes each example came
    from.
  Returns:
    Structural similarity computed per example, shape [batch, ...].
  """
  if volume_max is None:
    volume_max = torch.ones(logits.shape[0], device=DEVICE)
  # NOTE(dsuo): `volume_max` can be 0 if we have a padded batch, but this will
  # lead to NaN values in `ssim`.
  volume_max = torch.where(volume_max == 0,
                           torch.ones_like(volume_max),
                           volume_max)
  if mean is None:
    mean = torch.zeros(logits.shape[0], device=DEVICE)
  if std is None:
    std = torch.ones(logits.shape[0], device=DEVICE)
  mean = mean.view((-1,) + (1,) * (len(logits.shape) - 1))
  std = std.view((-1,) + (1,) * (len(logits.shape) - 1))
  logits = logits * std + mean
  targets = targets * std + mean
  ssims = functorch.vmap(structural_similarity)(logits, targets, volume_max)
  return ssims