code: def suggest_checkpoint_location(
      traces: List[LayerMemoryTrace],
      nb_checkpoints: int,
      num_skipped_layers: int,
  ) -> SuggestedCheckpoints:
      """
      Given a trace of a model, collected with or without checkpoint,
      return the best places to insert a reset of activation memory.
  
      The names of the returned modules are the boundaries of the
      suggested checkpoint_wrapper wrappings
      """
  
      # From the traces, extract how much activation memory
      # is generated during the forward pass, layer by layer
      visited = set()
      modules, allocations = [], []
      for t in traces:
          if t.is_forward:
              name = t.module_name
              memory = t.event.memory_activations
              if name not in visited:
                  visited.add(name)
                  modules.append(name)
                  allocations.append(memory)
  
      # remove the stem part
      modules = modules[num_skipped_layers:]
      allocations = allocations[num_skipped_layers:]
  
      # Compute the best positions to reset the memory
      max_memory, reset_indices = find_best_reset_points(
          allocations, nb_checkpoints=nb_checkpoints
      )
  
      # Then map it back to module names
      return SuggestedCheckpoints(
          max_memory=max_memory,
          split_modules=[modules[i] for i in reset_indices],
          all_modules=modules,
      )
  
imports: from contextlib import contextmanager
  from dataclasses import dataclass
  from enum import auto, Enum
  from functools import lru_cache
  from typing import Dict, List, NamedTuple, Tuple, Union
  import numpy
  import torch
  import torch.nn
  from fairscale.nn import FullyShardedDataParallel
  from vissl.utils.visualize import matplotlib_figure_to_image
label_desc: Use image node_id 2 for calling a global function with example usage: suggest_checkpoint_location(traces, nb_checkpoints, num_skipped_layers) and returns: SuggestedCheckpoints