Function leave_only_N

dali/test/python/reader/test_index.py:130–136

(src, dst, n)

    idx_file = os.path.join(idx_files_dir.name, recordio_idx)

    def leave_only_N(src, dst, n):
        with open(src, "r") as tmp_f:
            with open(dst, "w") as f:
                for i, x in enumerate(tmp_f):
                    if i == n:
                        break
                    f.write(x)

    leave_only_N(tfrecord_idx, idx_file, num_samples)