code: def _layer_memory_tracking_worker(
      gpu_id: int, sync_file: str, world_size: int
  ):
      init_distributed_on_file(
          world_size=world_size, gpu_id=gpu_id, sync_file=sync_file
      )
      torch.manual_seed(0)
      torch.backends.cudnn.deterministic = True
      torch.manual_seed(gpu_id)
  
      batch_size = 16
      fake_inputs = torch.randn(size=(batch_size, 10)).cuda(gpu_id)
      fake_targets = torch.randn(size=(batch_size, 10)).cuda(gpu_id)
      fake_criterion = nn.MSELoss()
  
      torch.manual_seed(0)
      torch.cuda.manual_seed(0)
  
      # Create a global group and a tracker around it
      group = dist.new_group()
      group = ProcessGroupTracker(group)
  
      # Create a simple model
      model = nn.Sequential(
          nn.Linear(10, 10).cuda(gpu_id),
          nn.ReLU(),
          FullyShardedDataParallel(
              nn.Linear(10, 10).cuda(gpu_id),
              flatten_parameters=False,
              process_group=group,
          ),
          nn.ReLU(),
          FullyShardedDataParallel(
              nn.Linear(10, 10).cuda(gpu_id),
              flatten_parameters=True,
              process_group=group,
          ),
      )
      model = model.cuda(gpu_id)
      model = FullyShardedDataParallel(
          model, flatten_parameters=False, process_group=group
      )
  
      # Setup the tracking of the model
      tracker = LayerwiseMemoryTracker()
      tracker.monitor(model)
  
      # Fake forward / backward pass
      fake_criterion(model(fake_inputs), fake_targets).backward()
  
      # Collect results of all gathers (the feature specific to FSDP)
      tracker.stop()
      all_gathered_traces = [
          (t.module_name, t.all_gathered, t.cumul_all_gathered)
          for t in tracker.memory_traces
          if t.all_gathered > 0
      ]
      assert all_gathered_traces == [
          ("_fsdp_wrapped_module._fpw_module.0", 440, 440),
          (
              "_fsdp_wrapped_module._fpw_module.2._fsdp_wrapped_module._fpw_module",
              440,
              880,
          ),
          (
              "_fsdp_wrapped_module._fpw_module.4._fsdp_wrapped_module._fpw_module",
              440,
              880,
          ),
          (
              "_fsdp_wrapped_module._fpw_module.4._fsdp_wrapped_module._fpw_module",
              440,
              0,
          ),
          (
              "_fsdp_wrapped_module._fpw_module.2._fsdp_wrapped_module._fpw_module",
              440,
              0,
          ),
      ], f"Expected {all_gathered_traces}"
  
file_path: vissl/tests/test_layer_memory_tracking.py
imports: import unittest
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  import torchvision.models
  from fairscale.nn import FullyShardedDataParallel
  from vissl.utils.layer_memory_tracking import find_best_reset_points, LayerwiseMemoryTracker, ProcessGroupTracker
  from vissl.utils.test_utils import gpu_test, init_distributed_on_file, with_temp_files, with_timing
label_desc: Use image node_id 2 for calling the TestLayerMemoryTracking obj's underlying member method code with example usage: obj._layer_memory_tracking_worker(gpu_id, sync_file, world_size) without return types