code: def profile_model(args, model, inputs, rank):
      with profile(
          activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]
      ) as prof:
          for i in range(args.repeat):
              with record_function("Forward"):
                  outputs = model(*inputs)
                  loss = reduce_to_scalar_loss(outputs)
              with record_function("Backward"):
                  loss.backward()
      if rank == 0:
          prof.export_chrome_trace(args.trace_file)
  
file_path: pytorch/benchmarks/dynamo/distributed.py
imports: import argparse
  import logging
  import os
  from functools import partial
  import torch
  import torch._dynamo
  import torch.utils._pytree
  from torch._dynamo.testing import reduce_to_scalar_loss
  from torch.nn.parallel import DistributedDataParallel
  from torch.profiler import profile, ProfilerActivity, record_function
label_desc: Use image node_id 2 for calling a global function with example usage: profile_model(args, model, inputs, rank) without return types