code: def test_memory_tracking_nlp_model():
      """
      Check that we can collect memory traces of a realistic model
      outside of the context of distributed training (DDP or FSDP)
      """
  
      BACH_SIZE = 10
      INPUT_DIM = 16
      model = GPT2(
          embed_dim=256,
          num_heads=2,
          num_layers=6,
          num_positions=INPUT_DIM * INPUT_DIM,
          num_vocab=512,
          num_classes=2,
      ).cuda()
      tracker = LayerwiseMemoryTracker()
      tracker.monitor(model)
      input_tensor = torch.randint(10, (BACH_SIZE, INPUT_DIM)).cuda()
      output = model(input_tensor)
      output.sum().backward()
  
      assert (
          len(tracker.memory_traces) > 0
      ), "failed to collected memory traces"
      assert (
          len(tracker.forward_traces) > 0
      ), "failed to collect forward memory traces"
      assert (
          len(tracker.backward_traces) > 0
      ), "failed to collect backward memory traces"
      assert tracker.summary.total_activation_allocations == 12462080
  
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 2 for calling a global function with example usage: test_memory_tracking_nlp_model() without return types