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

RUN pip install runpod transformers

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

RUN python -c 'from transformers import pipeline; import torch; pipe = pipeline("text-generation", model="rwitz/go-bruins-v2",device=0,torch_dtype=torch.bfloat16)'

ADD handler.py .

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