How to predict with 1000s of dataloaders?

[HadiSDev]

I have a case where I have 2K + dataloaders that each hold a dataset with a dataframe.

The model is a pretrained Bert model, and I want to create a trainer that can take all these dataloaders and call predict on them.

Each output should be seperate, meaning I cannot mix between the different dataloaders. I need to save the predictions individually for each dataset.

        trainer = pl.Trainer(
            accelerator="gpu",
            precision="bf16-mixed",
            strategy=("ddp" if torch.cuda.device_count() > 1 else "auto"),
            callbacks=writer,
        )
        trainer.predict(model, dataloaders=all_dataloaders)

The problem I am having is I get random OOM from cuda after 20th, 50th, 100th or 200th dataloader, so I am never able to actually finish the predict step. It seems like the predict function was never made to handle so many dataloaders. I tried making it work with 1 or 2 GPUs (L4) but still no success :(

The only solution I am left with is running it with 1 GPU with dataloaders in batches instead, save predictions, and send the next batch of dataloaders, but this does not really work with multi-gpu setup.

Any ideas what I am doing wrong?