gui11aume/pyro-gpu repository overview

# pyro-gpu:1.9.0
FROM nvidia/cuda:12.3.1-devel-ubuntu22.04

RUN apt-get update -y
RUN apt-get install -y wget

# Install Python and PIP.
RUN apt-get install -y python3 python3-pip && apt-get clean && rm -rf /var/lib/apt/lists/*
RUN ln /usr/bin/python3 /usr/bin/python

# Install PyTorch, Pyro Lightning and Deepspeed
RUN pip3 install deepspeed==0.13.1 lightning==2.2.0 pyro-ppl==1.9.0 torch==2.2.0