multi-gpu inference via ddp for pytorch-lightning/lightning-ai models

pl_multigpu_infer.py

  import pytorch_lightning as pl
  ...
  # override this method on pytorch-lightning model
  def on_predict_epoch_end(self, results):
      # gather all results onto each device
      # find created world_size from pl.trainer
      results = all_gather(results[0], WORLD_SIZE, self._device)
      # concatenate on the cpu
      results = torch.concat([x.cpu() for x in results], dim=1)
      # output will not preserve input order. 
      # suggest outputing index in predict step and sort by index value here.
      ...