code: def test_memory_tracking_traces():
      """
      Minimal test case to check that we can collect memory traces
      outside of the context of distributed training (DDP or FSDP)
      """
  
      # Create a model with a hierarchy of modules
      torch.manual_seed(0)
      model = nn.Sequential(
          nn.Sequential(
              nn.Conv2d(
                  3, 64, kernel_size=(3, 3), padding=(1, 1), bias=False
              ),
              nn.BatchNorm2d(64),
              nn.ReLU(inplace=False),
              nn.AdaptiveAvgPool2d(output_size=(1, 1)),
          ),
          nn.Flatten(start_dim=1),
          nn.Sequential(nn.Linear(64, 2), nn.ReLU(inplace=True)),
      ).cuda()
  
      # Track a fake forward / backward
      tracker = LayerwiseMemoryTracker()
      tracker.monitor(model)
      x = torch.randn(size=(2, 3, 224, 224)).cuda()
      target = torch.LongTensor([0, 1]).cuda()
      criterion = nn.CrossEntropyLoss()
      criterion(model(x), target).backward()
  
      # Verify that only leaf modules are tracked and that the order
      # of the traces is consistent with backward/forward
      tracked_names = [t.module_name for t in tracker.memory_traces]
      expected_names = ["0.0", "0.1", "0.2", "0.3", "1", "2.0", "2.1"]
      assert set(expected_names) == set(tracked_names)
      assert tracked_names == (expected_names + expected_names[::-1])
  
      # Verify that memory tracking for ReLU is sound
      assert (
          2 * 64 * 224 * 224 * 4
          == tracker.forward_traces[2].event.memory_activations
      ), "ReLU(inplace=False) should allocate activations"
      assert (
          0 == tracker.forward_traces[6].event.memory_activations
      ), "ReLU(inplace=True) should NOT allocate activations"
  
      # Verify that overall memory tracking is sound
      summary = tracker.summary
      assert (
          summary.total_forward_allocations
          >= summary.total_activation_allocations
      )
  
      # Verify that the identification of top memory activation producer works:
      # these are the first layers, all allocating (2, 64, 224, 224) feature maps
      top_act_producers = summary.top_forward_activation_producers[:3]
      assert "0.0" == top_act_producers[0].module_name
      assert "0.1" == top_act_producers[1].module_name
      assert "0.2" == top_act_producers[2].module_name
      assert 3 * 3 * 64 * 3 * 4 == top_act_producers[0].module_params
      assert 64 * 2 * 4 == top_act_producers[1].module_params
      assert 0 == top_act_producers[2].module_params
      for trace in top_act_producers:
          assert (
              2 * 64 * 224 * 224 * 4 == trace.event.memory_activations
          )
  
file_path: fairscale/tests/experimental/tooling/test_layer_memory_tracker.py
imports: from typing import Tuple
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  from torch.nn.parallel import DistributedDataParallel
  from fairscale.experimental.tooling.layer_memory_tracker import LayerwiseMemoryTracker, ProcessGroupTracker, find_best_reset_points
  from fairscale.fair_dev.testing.testing import GPT2, dist_init, skip_if_no_cuda, skip_if_single_gpu, temp_files_ctx
  from fairscale.nn import FullyShardedDataParallel
label_desc: Use image node_id 1 for calling a global function with example usage: test_memory_tracking_traces() without return types