FROM pytorch/pytorch:1.12.0-cuda11.3-cudnn8-runtime

ADD ./main.py /workspace/main.py

ENTRYPOINT ["python", "/workspace/main.py"]