code: def get_partitioned_activations_for_backward(
      args, inputs, contiguous_checkpoint
  ):
      global contiguous_size_buffers, size_offsets
  
      new_args = []
      num_non_fp_tensors = 0
  
      for arg_index, (arg, inp) in enumerate(zip(args, inputs)):
          size = (
              torch.tensor(arg.size()) if torch.is_tensor(arg) else None
          )
          if not is_activation_to_checkpoint(arg):
              new_args.append(arg)
              new_args.append(size)
              num_non_fp_tensors += 1
              continue
  
          arg.data = torch.empty([], device=arg.device).data
          arg.saved_data = inp.data
  
          new_args.append(arg)
          i = arg_index - num_non_fp_tensors
  
          if contiguous_checkpoint:
              numel = size.numel()
              if i >= len(contiguous_size_buffers):
                  tmp = torch.tensor(())
                  contiguous_size_buffers.append(
                      tmp.new_empty(
                          [numel * num_layers],
                          dtype=size.dtype,
                          device=size.device,
                      )
                  )
                  size_offsets.append(0)
              elif contiguous_size_buffers[i] is None:
                  tmp = torch.tensor(())
                  contiguous_size_buffers[i] = tmp.new_empty(
                      [numel * num_layers],
                      dtype=size.dtype,
                      device=size.device,
                  )
                  size_offsets[i] = 0
  
              contiguous_size = (
                  contiguous_size_buffers[i]
                  .narrow(0, size_offsets[i], numel)
                  .data.copy_(size.data)
              )
              contiguous_size = contiguous_size.view_as(size)
              size_offsets[i] = size_offsets[i] + numel
              new_args.append(contiguous_size)
          else:
              new_args.append(size)
  
      return new_args
  
imports: import copy
  import torch
  import contextlib
  from deepspeed import comm
  import weakref
  import mmap
  from torch import _C
  from deepspeed.runtime.config import DeepSpeedConfig
  from deepspeed.utils import logger
  from deepspeed.runtime.utils import copy_to_device, move_to_device, see_memory_usage, bwc_tensor_model_parallel_rank
  from deepspeed.utils.timer import SynchronizedWallClockTimer, FORWARD_GLOBAL_TIMER
  from deepspeed.accelerator import get_accelerator
label_desc: Use image node_id 13 for calling a global function with example usage: get_partitioned_activations_for_backward(args, inputs, contiguous_checkpoint) and returns: new_args