code: def main():
      parser = argparse.ArgumentParser("Benchmarking MEVO")
  
      parser.add_argument(
          "--dtype", type=str, choices=["fp16", "fp32"], default="fp16"
      )
      parser.add_argument(
          "--grad",
          type=str,
          choices=["grad", "no_grad"],
          default="grad",
      )
      parser.add_argument(
          "--fwd_bwd", action="store_true", default=False
      )
      args = parser.parse_args()
  
      repeats = 9
      results = {}
      results["peak cached"] = {}
      results["durations"] = {}
      for shape in SHAPES:
          name = shape[0]
          results["peak cached"][name] = {}
          results["durations"][name] = {}
          dtype = (
              torch.float32 if args.dtype == "fp32" else torch.float16
          )
          # Use cpu memory to ensure we always start with an empty GPU
          data = get_data(shape[1:], dtype, "cpu")
          for kernel in KERNELS:
              k_name = kernel.__name__
              no_grad = args.grad
              print(
                  f"Running {k_name} with {name} {dtype} {no_grad} data"
              )
              peak_mem, durations = run_on_gpu(
                  kernel,
                  data,
                  repeats,
                  no_grad == "no_grad",
                  args.fwd_bwd,
              )
              results["peak cached"][name][k_name] = peak_mem
              results["durations"][name][k_name] = durations
      pprint(results)
  
file_path: fairscale/benchmarks/experimental/benchmark_mevo.py
imports: import argparse
  import contextlib
  from pprint import pprint
  from statistics import mean
  import time
  import torch
  from torch import nn
  from torch.cuda import Event
  from fairscale.experimental.nn import MEVO, BaselineSoftmaxNllLoss
  from fairscale.experimental.nn.mevo import get_data
label_desc: Use image node_id 2 for calling a global function with example usage: main() without return types