From f9c9fbe0186fed700fa9180ae5960926d38f3613 Mon Sep 17 00:00:00 2001
From: Zach Koopmans <zkoopmans@google.com>
Date: Fri, 23 Jun 2023 10:34:37 -0700
Subject: [PATCH] Add initial pytorch container definition.

Add initial pytorch container image for pytorch benchmarks.

PiperOrigin-RevId: 542897253
---
 images/benchmarks/pytorch/Dockerfile.x86_64 | 23 +++++++++++++++++++++
 1 file changed, 23 insertions(+)
 create mode 100644 images/benchmarks/pytorch/Dockerfile.x86_64

diff --git a/images/benchmarks/pytorch/Dockerfile.x86_64 b/images/benchmarks/pytorch/Dockerfile.x86_64
new file mode 100644
index 000000000..38734f3ef
--- /dev/null
+++ b/images/benchmarks/pytorch/Dockerfile.x86_64
@@ -0,0 +1,23 @@
+FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel
+
+RUN apt-get update
+RUN DEBIAN_FRONTEND=noninteractive apt-get install -y git wget \
+    libgl1-mesa-glx libglib2.0-0
+
+RUN git clone https://github.com/pytorch/benchmark.git \
+    && cd /workspace/benchmark \
+    && git reset --hard 6d84522ec83997beefa4bb6e5154453b5fc653f7
+
+# For some reason, installing the suite's requirements.txt does not correctly
+# install "boto3", so install it explicitly.
+RUN pip install boto3 && pip install -r /workspace/benchmark/requirements.txt
+
+RUN cd /workspace/benchmark && python install.py \
+    DALLE2_pytorch \
+    LearningToPaint \
+    Super_SloMo
+
+RUN cd /workspace/benchmark/torchbenchmark/models && pip install \
+    -r DALLE2_pytorch/requirements.txt \
+    -r LearningToPaint/requirements.txt \
+    -r Super_SloMo/requirements.txt