def get_bond_features_shape(train_set):
    """Get the shape of the atom features from a MoleculeACE.GNN.dataloaders.*_graph_dataset"""
    bond_features = 13  # this is the standard size
    for item in train_set:
        single_graph = item[1]
        bond_features = single_graph.edata['e'].shape[1]
        break

    return bond_features