What we can do here is a in_memory_cache over image_files_dp. Then, we could add an API to convert a IterDataPipe to a lazy loaded MapDataPipe to represent the LazyDict.
If we use LazyDict here, I have concern that image_files_map would be missing from the DataLoader graph.
cc: @VitalyFedyunin

Besides, I think the DataLoader would complain this datapipe graph in the second epoch because image_files_dp is never used after the first epoch then Demux would also be non-serializable same as the comment I made in the other PR.

So, a fix from Demux is not avoidable.

This seems like a good thing to test in general. What should a test look like. Is something like

for _ in dataset.cycle(2):
    pass

enough? If yes, my proposal passes this test.

I mean if we put the dataset (datapipes) into DataLoader, the second epoch of DataLoader would break.

So something like

data_loader = DataLoader2(dataset)

for epoch in range(2):
    for sample in data_loader:
        pass

?

Yeah. I have asked Kevin to fix such issue in demux.

My proposal still works. I've pushed the test I'm running against. There are multiple failures for other datasets, but cub200 is not one of them.

A second thought. Could we simply filter image_files_dp from archive_dp here and create the image_files_map dictionary?

Then, we can do demux over archive_dp again and drop data in image_files_dp.

So basically splitting of image_files_dp from the graph?

Yeah. Then, we can materialize the data from it like a meta-datapipe.