def get_node_map(model: torch.fx.GraphModule) -> Dict[str, Node]:
    """
    Return a dictionary of node name to node
    """
    node_map = {}
    for node in model.graph.nodes:
        node_map[node.name] = node
    return node_map

    Return a dictionary of node name to node
    
coremltools/test/optimize/torch/quantization/test_coreml_quantizer.py