FROM nvcr.io/nvidia/pytorch:24.01-py3

WORKDIR /app

RUN pip install --no-cache-dir datasets huggingface_hub

COPY src/ /app/src/
COPY experiments/ /app/experiments/
COPY train_gpu.py /app/train_gpu.py

CMD ["python", "train_gpu.py"]