FROM nvcr.io/nvidia/driver:550.127.08-ubuntu22.04

RUN apt-get update && apt-get install -y python3 python3-pip
WORKDIR /app

COPY . .
RUN pip install -r requirements.txt

CMD ["python3", "train.py"]