class_docstr: Similar to the torch version, but support non-Tensor outputs.
  
  The caller is expected to provide a dict (*parent_ctx_dict*) that will hold
  the non-Tensor outputs. These should be combined with the Tensor *outputs*
  by calling :func:`unpack_non_tensors`.
class_docstr_tok: ["Similar","to","the","torch","version",",","but","support","non-Tensor","outputs",".","The","caller","is","expected","to","provide","a","dict","(","*","parent_ctx_dict","*",")","that","will","hold","the","non-Tensor","outputs",".","These","should","be","combined","with","the","Tensor","*","outputs","*","by","calling",":","func",":","`","unpack_non_tensors","`","."]
code: def forward(  # type: ignore
      ctx: Any,
      dummy_tensor_requires_grad: torch.Tensor,
      run_function: Any,
      parent_ctx_dict: Dict[str, Any],
      kwarg_keys: Tuple[str, ...],
      *args: Any,
      **kwargs: Any,
  ) -> Any:
      torch_checkpoint.check_backward_validity(args)
  
      ctx.run_function = run_function
      ctx.kwarg_keys = kwarg_keys
      ctx.fwd_rng_state = get_rng_state()
      ctx.had_autocast_in_fwd = is_autocast_enabled()
  
      tensor_inputs, packed_non_tensor_inputs = split_non_tensors(args)
      if parent_ctx_dict["offload"]:
          ctx.fwd_device = tuple(x.device for x in tensor_inputs)
          ctx.grad_requirements = tuple(
              x.requires_grad for x in tensor_inputs
          )
          tensor_inputs = tuple(
              x.to("cpu", non_blocking=True) for x in tensor_inputs
          )
      else:
          ctx.fwd_device, ctx.grad_requirements = None, None
  
      ctx.save_for_backward(*tensor_inputs)
      ctx.packed_non_tensor_inputs = packed_non_tensor_inputs
  
      with torch.no_grad(), enable_checkpointing():
          unpacked_args, unpacked_kwargs = unpack_kwargs(
              kwarg_keys, args
          )
          outputs = run_function(*unpacked_args, **unpacked_kwargs)
          the_module = unpacked_args[0]
  
      # Because we run with torch.no_grad(), we can't actually access
      # outputs.requires_grad. Instead, we manually compute it by
      # checking if either the input or the module needs grads
      parameters = list(the_module.parameters())
  
      # If the module is wrapped by FlattenParamsWrapper, then the
      # parameters would have been deleted. If so, we need to access
      # the views into the flattened parameters.
      if hasattr(the_module, "_unflattened_param_views"):
          parameters += the_module._unflattened_param_views
  
      output_requires_grad = any(
          param.requires_grad for param in parameters
      ) or any(x.requires_grad for x in tensor_inputs)
      parent_ctx_dict["output_requires_grad"] = output_requires_grad
  
      if not isinstance(outputs, torch.Tensor):
          # Autograd Functions don't like non-Tensor outputs. We can split the
          # non-Tensor and Tensor outputs, returning the former by reference
          # through *parent_ctx_dict* and returning the latter directly.
          outputs, packed_non_tensor_outputs = split_non_tensors(
              outputs
          )
          parent_ctx_dict[
              "packed_non_tensor_outputs"
          ] = packed_non_tensor_outputs
  
      return outputs
  
imports: from contextlib import contextmanager
  from dataclasses import dataclass
  import functools
  import threading
  from typing import Any, Dict, Generator, Optional, Tuple
  import weakref
  import torch
  from torch import Tensor
  import torch.nn
  import torch.utils.checkpoint
  from fairscale.internal.containers import pack_kwargs, split_non_tensors, unpack_kwargs, unpack_non_tensors
  from .checkpoint_utils import patch_batchnorm
label_desc: Use image node_id 1 for calling the CheckpointFunction obj's underlying member method code with example usage: obj.forward(ctx, dummy_tensor_requires_grad, run_function, parent_ctx_dict, kwarg_keys) and returns: outputs