I have a use case where my data is stored in multiple files where each files can have multiple batches. I am using Pytorch Lightning and I currently have multiple GPUs(Ranks) and I use a Dataloader with multiple workers for each Rank. Since DMP requires even number of batches per Rank even in inference mode, I am creating dummy batches on my workers after they are exhausted with their real valid data.
This results in two issues:

1. For a single rank, since different workers have different number of real batches, at one point many workers for this rank yield dummy matches, while some workers yield real, so there is wastage of computation.
2. For different ranks, it's hard to communicate when all ranks have exhausted the data, so I set a limit_val_batches etc, but this is not ideal.

Does anyone have any advice or recommendations for this use case?