code: def benchmark(
      name: str,
      model: torch.nn.Module,
      warmup_inputs: List[KeyedJaggedTensor],
      bench_inputs: List[KeyedJaggedTensor],
      prof_inputs: List[KeyedJaggedTensor],
      batch_size: int,
      world_size: int,
      output_dir: str,
      num_benchmarks: int,
  ) -> BenchmarkResult:
      model.training = False
      max_mem_allocated: List[int] = []
      logger.info(f" BENCHMARK_MODEL[{name}]:\n{model}")
  
      for _input in warmup_inputs:
          model(_input)
  
      # Reset memory for measurement
      for di in range(torch.cuda.device_count()):
          torch.cuda.reset_max_memory_allocated(
              torch.device(f"cuda:{di}")
          )
  
      # Measure time taken for batches in bench_inputs
      start = [
          torch.cuda.Event(enable_timing=True)
          for _ in range(num_benchmarks)
      ]
      end = [
          torch.cuda.Event(enable_timing=True)
          for _ in range(num_benchmarks)
      ]
  
      gc.disable()
  
      for i in range(num_benchmarks):
          start[i].record()
          for _input in bench_inputs:
              model(_input)
          end[i].record()
          gc.collect()
  
      for di in range(torch.cuda.device_count()):
          torch.cuda.synchronize(torch.device(f"cuda:{di}"))
  
      # TODO: First Benchmark Run for Eager Mode produces outlier
      # Start counting after first as workaround for standard deviation
      elapsed_time = torch.tensor(
          [si.elapsed_time(ei) for si, ei in zip(start[1:], end[1:])]
      )
  
      total_duration_sec = (
          elapsed_time.mean().item() * 1e-3
      )  # time in seconds
      std_sec = elapsed_time.std().item() * 1e-3  # time in seconds
      gc.enable()
  
      for di in range(world_size):
          b = torch.cuda.max_memory_allocated(
              torch.device(f"cuda:{di}")
          )
          max_mem_allocated.append(b // 1024 // 1024)
  
      # pyre-ignore[2]
      def trace_handler(prof) -> None:
          total_average = prof.profiler.total_average()
          logger.info(f" TOTAL_AVERAGE:\n{name}\n{total_average}")
          dir_path: str = output_dir
          trace_file: str = f"{dir_path}/trace-{name}.json"
          stacks_cpu_file = f"{dir_path}/stacks-cpu-{name}.stacks"
          stacks_cuda_file = f"{dir_path}/stacks-cuda-{name}.stacks"
          logger.info(f" PROFILE[{name}].chrome_trace:{trace_file}")
  
          prof.export_chrome_trace(trace_file)
          prof.export_stacks(stacks_cpu_file, "self_cpu_time_total")
          prof.export_stacks(stacks_cuda_file, "self_cuda_time_total")
  
      # - git clone https://github.com/brendangregg/FlameGraph
      # - cd FlameGraph
      # - ./flamegraph.pl --title "CPU time" --countname "us." profiler.stacks > perf_viz.svg
  
      with torch.profiler.profile(
          activities=[
              torch.profiler.ProfilerActivity.CPU,
              torch.profiler.ProfilerActivity.CUDA,
          ],
          record_shapes=True,
          profile_memory=True,
          with_stack=True,
          with_flops=True,
          with_modules=True,
          on_trace_ready=trace_handler,
      ) as p:
          for _input in prof_inputs:
              with record_function("## forward ##"):
                  model(_input)
                  p.step()
          for di in range(torch.cuda.device_count()):
              torch.cuda.synchronize(torch.device(f"cuda:{di}"))
  
      return BenchmarkResult(
          short_name=name,
          total_duration_sec=total_duration_sec,
          std_sec=std_sec,
          max_mem_allocated=max_mem_allocated,
      )
  
imports: import argparse
  import gc
  import logging
  import os
  import time
  from dataclasses import dataclass
  from enum import Enum
  from typing import Dict, List, Tuple
  import torch
  from torch.autograd.profiler import record_function
  from torchrec.distributed.embedding_types import EmbeddingComputeKernel, ShardingType
  from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
  from torchrec.distributed.planner.enumerators import EmbeddingEnumerator
  from torchrec.distributed.planner.shard_estimators import EmbeddingPerfEstimator, EmbeddingStorageEstimator
  from torchrec.distributed.shard import _shard_modules
  from torchrec.distributed.test_utils.infer_utils import TestQuantEBCSharder
  from torchrec.distributed.test_utils.test_model import ModelInput
  from torchrec.distributed.types import DataType, ShardingEnv
  from torchrec.fx import symbolic_trace
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
  from torchrec.quant.embedding_modules import EmbeddingBagCollection
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, KeyedTensor
label_desc: Use image node_id 3 for calling a global function with example usage: benchmark(name, model, warmup_inputs, bench_inputs, prof_inputs, batch_size, world_size, output_dir, num_benchmarks) and returns: BenchmarkResult