code: def custom_backward(
      output: torch.Tensor, grad_output: Optional[torch.Tensor]
  ) -> None:
      """Directly call C++ autograd engine.
  
      To make the `free_output_tensor` optimization work, the C++ autograd engine must be called
      directly, bypassing PyTorch's `torch.autograd.backward`. PyTorch's `backward` checks that the
      output and grad have the same shape, while C++ `backward` does not.
      """
      assert (
          output.numel() == 1
      ), "output should be pseudo-freed in schedule, to optimize memory consumption"
      assert isinstance(output, torch.Tensor), "output == {}.".format(
          type(output).__name__
      )
      assert isinstance(
          grad_output, (torch.Tensor, type(None))
      ), "grad_outptu == {}.".format(type(grad_output).__name__)
  
      # Handle scalar output
      if grad_output is None:
          assert (
              output.numel() == 1
          ), "Implicit grad requires scalar output."
          grad_output = torch.ones_like(
              output, memory_format=torch.preserve_format
          )
  
      # Call C++ engine [ see torch/csrc/autograd/python_engine.cpp ]
      Variable._execution_engine.run_backward(
          tensors=(output,),
          grad_tensors=(grad_output,),
          keep_graph=False,
          create_graph=False,
          inputs=(),
          allow_unreachable=True,
          accumulate_grad=True,
      )
  
file_path: apex/apex/transformer/pipeline_parallel/schedules/common.py
imports: from typing import Any, Callable, Dict, List, Tuple, Union, Optional, Sequence
  import torch
  from torch.autograd.variable import Variable
  from apex.normalization.fused_layer_norm import FusedLayerNorm
  from apex.transformer import parallel_state
  from apex.transformer.enums import ModelType
  from apex.transformer.pipeline_parallel.p2p_communication import FutureTensor
  from apex.transformer.pipeline_parallel.utils import get_num_microbatches
  from apex.transformer.pipeline_parallel.utils import listify_model
  from apex.transformer.pipeline_parallel.utils import unwrap_model
  from apex.transformer.pipeline_parallel.utils import get_model_type
  from apex.transformer.tensor_parallel.layers import set_defaults_if_not_set_tensor_model_parallel_attributes
  from apex.transformer.log_util import get_transformer_logger
label_desc: Use image node_id 5 for calling a global function with example usage: custom_backward(output, grad_output) without return types