When we use lightning in multigpu setting with strategy ddp,
predictions = trainer.predict(model, dataloaders=predict_dataloader)
the predictions are from a single gpu. From Combining loss, predictions from multi gpus setting in pytorch lightning it looks like we can use self.all_gather to collect predictions from all gpus. However, how can I align the collected predictions with the dataset used by predict_dataloader?
I know torchmetrics can be used to help with metrics, but I'm not computing metrics. I want to collect all predictions and be able to process each prediction with it's corresponding input. And I want to do so using strategy ddp with multiple gpus.
