def batch_graph_data_dict(batched_data_dict):
    """
    After running dataset.batch() on data_dict representation of GraphTuple, correct the batch dimensions.

    Args:
        batched_data_dict: dict(
            nodes[num_graphs, n_node_per_graph, F_nodes],
            edges[num_graphs, n_edge_per_graph, F_edges],
            senders[num_graphs, n_edge_per_graph],
            receivers[num_graphs, n_edge_per_graph],
            globals[num_graphs, 1, F_globals],
            n_node[num_graphs, 1],
            n_edge[num_graphs, 1])

    Returns:
        batched_data_dict representing a batched GraphTuple:
        dict(
            nodes[num_graphs, n_node_per_graph, F_nodes],
            edges[num_graphs, n_edge_per_graph, F_edges],
            senders[num_graphs, n_edge_per_graph],
            receivers[num_graphs, n_edge_per_graph],
            globals[num_graphs, F_globals],
            n_node[num_graphs],
            n_edge[num_graphs])
    """
    if "globals" in batched_data_dict.keys():
        batched_data_dict["globals"] = batched_data_dict["globals"][:,0,:]
    if "n_node" in batched_data_dict.keys():
        batched_data_dict['n_node'] = batched_data_dict['n_node'][:,0]
    if "n_edge" in batched_data_dict.keys():
        batched_data_dict['n_edge'] = batched_data_dict['n_edge'][:,0]
    return batched_data_dict