code: def with_nvtx_range(
      layer: _ModelT,
      name: Optional[str] = None,
      *,
      forward_color: int = -1,
      backprop_color: int = -1,
  ) -> _ModelT:
      """Wraps any layer and marks the forward and backprop phases as
      NVTX ranges for CUDA profiling.
  
      By default, the name of the layer is used as the name of the range,
      followed by the name of the pass.
      """
      name = layer.name if name is None else name
  
      orig_forward = layer._func
      orig_init = layer.init
  
      def forward(
          model: Model, X: Any, is_train: bool
      ) -> Tuple[Any, Callable]:
          with use_nvtx_range(f"{name} forward", forward_color):
              layer_Y, layer_callback = orig_forward(
                  model, X, is_train=is_train
              )
  
          def backprop(dY: Any) -> Any:
              with use_nvtx_range(f"{name} backprop", backprop_color):
                  return layer_callback(dY)
  
          return layer_Y, backprop
  
      def init(_model: Model, X: Any, Y: Any) -> Model:
          if orig_init is not None:
              return orig_init(layer, X, Y)
          else:
              return layer
  
      layer.replace_callbacks(forward, init=init)
  
      return layer
  
imports: from typing import Any, Callable, Optional, Tuple, TypeVar
  from ..model import Model
  from ..util import use_nvtx_range
label_desc: Use image node_id 1 for calling a global function with example usage: with_nvtx_range(layer, name) and returns: layer, layer_Y,  backprop, orig_init, layer, layer_callback