Skip to content

Instantly share code, notes, and snippets.

@will-thompson-k
will-thompson-k / pl_multigpu_infer.py
Created August 15, 2023 15:09
multi-gpu inference via ddp for pytorch-lightning/lightning-ai models
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
will-thompson-k / nvidia_install.sh
Created September 21, 2023 16:17
if nvidia-smi fails to work
sudo apt-get remove --purge '^nvidia-.*'
sudo ubuntu-drivers autoinstall
sudo reboot