Dockerfile

FROM runpod/pytorch:2.0.1-py3.10-cuda11.8.0-devel

WORKDIR /workspace

RUN pip install runpod transformers

RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

ADD handler.py .

CMD [ "python", "-u", "/handler.py" ]