using tensorflow Dataset api to do data_argument

profiler.py

    import cProfile, pstats
    from pstats import SortKey

    sortby = SortKey.CUMULATIVE
    with cProfile.Profile() as pr:


        # your code here

        ps = pstats.Stats(pr).sort_stats(sortby)
        ps.print_stats()