will-thompson-k

  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. 

will-thompson-k

sudo apt-get remove --purge '^nvidia-.*'
sudo ubuntu-drivers autoinstall
sudo reboot