Proper way of making predictions with multi-GPUs

[marcmk6]

Hi,

I implemented predict_step and tried to make predictions on multiple GPUs under DDP (e.g. predictions = trainer.predict(model, dataloaders=loader)), but in the end predictions contains only 1/N of total data, where N is the number of GPUs.

After observing this, I tried to override the hook on_predict_epoch_end, which is supposed to get results from predict_step, but it receives empty lists.

I'm wondering what's correct way of doing this.

Thanks