def METHOD_NAME(model, graph):
  module_to_node = {}
  for name, module in model.named_children():
    # module_name -> node_id
    idx = int(name.split(TorchSymbol.MODULE_NAME_SEPERATOR)[-1])
    node = graph.get_node_by_idx(idx)
    module_to_node[name] = node
  return module_to_node