code: def _smoothness_loss_fn(
      img_gx: Tensor, img_gy: Tensor, val_gx: Tensor, val_gy: Tensor
  ):
      # ref: https://github.com/nianticlabs/monodepth2/blob/b676244e5a1ca55564eb5d16ab521a48f823af31/layers.py#L202
  
      torch._assert(
          img_gx.ndim >= 3,
          "smoothness_loss: `img_gx` must be at least 3-dimensional tensor of shape (..., C, H, W)",
      )
  
      torch._assert(
          img_gx.ndim == val_gx.ndim,
          "smoothness_loss: `img_gx` and `depth_gx` must have the same dimensionality, but got {} and {}".format(
              img_gx.ndim, val_gx.ndim
          ),
      )
  
      for idx in range(img_gx.ndim):
          torch._assert(
              (
                  img_gx.shape[idx] == val_gx.shape[idx]
                  or (img_gx.shape[idx] == 1 or val_gx.shape[idx] == 1)
              ),
              "smoothness_loss: `img_gx` and `depth_gx` must have either the same shape or broadcastable shape, but got {} and {}".format(
                  img_gx.shape, val_gx.shape
              ),
          )
  
      # -3 is channel dimension
      weights_x = torch.exp(
          -torch.mean(torch.abs(val_gx), axis=-3, keepdim=True)
      )
      weights_y = torch.exp(
          -torch.mean(torch.abs(val_gy), axis=-3, keepdim=True)
      )
  
      smoothness_x = img_gx * weights_x
      smoothness_y = img_gy * weights_y
  
      smoothness = (
          torch.abs(smoothness_x) + torch.abs(smoothness_y)
      ).mean(axis=(-3, -2, -1))
      return smoothness
  
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 4 for calling a global function with example usage: _smoothness_loss_fn(img_gx, img_gy, val_gx, val_gy) and returns: smoothness