code: def track_time(graph, num_heads):
      device = utils.get_bench_device()
      graph = utils.get_graph(graph).to(device)
      score = (
          torch.randn((graph.num_edges(), num_heads))
          .requires_grad_(True)
          .float()
          .to(device)
      )
  
      # dry run
      for i in range(3):
          y = dgl.ops.edge_softmax(graph, score)
  
      # timing
      with utils.Timer(device) as t:
          for i in range(100):
              y = dgl.ops.edge_softmax(graph, score)
  
      return t.elapsed_secs / 100
  
imports: import time
  import dgl
  import torch
  from ..None import utils
label_desc: Use image node_id 1 for calling a global function with example usage: track_time(graph, num_heads) and returns: unknown