Optimizing dictionary data structures fails when using a partially initialized function

[enrico-stauss]

When passing fn=partial(my_func, arg1=arg1) to optimize, the line linked below will yield only the keys of the dictionary that my_funcreturns. I avoided the issue by not using partial for my usecase but I`d consider this a bug.

Kind regards :)

https://github.com/Lightning-AI/litdata/blob/a09de861eadcb8389d0bc3065a901113c9e72a5f/src/litdata/processing/functions.py#L138