FROM unsloth/unsloth:latest

WORKDIR /app

# Additional deps
RUN pip install --no-cache-dir einops jaxtyping

COPY train.py /app/train.py

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