code: def time_fwd_bwd(func, *args, **kwargs):
      time_f, time_b = benchmark_fwd_bwd(func, *args, **kwargs)
      return time_f[1].mean, time_b[1].mean
  
file_path: flash-attention/benchmarks/benchmark_flash_attention.py
imports: import pickle
  import math
  import torch
  import torch.nn
  import torch.nn.functional
  from einops import rearrange, repeat
  from flash_attn.utils.benchmark import benchmark_all, benchmark_forward, benchmark_backward
  from flash_attn.utils.benchmark import benchmark_fwd_bwd, benchmark_combined
  from flash_attn import flash_attn_qkvpacked_func
label_desc: Use image node_id 4 for calling a global function with example usage: time_fwd_bwd(func) and returns: time_f,  time_b