def get_graph(graph):
    """
    Get correct graph due to data parallel

    :param graph: network graph
    :type graph: torch.nn.Module
    :return: correct graph
    :rtype: torch.nn.Module
    """
    if hasattr(graph, 'module'):
        # fix DataParallel
        return graph.module
    else:
        return graph