From 565a1876ccfe65696f04938541c4f90ffe310064 Mon Sep 17 00:00:00 2001
From: Sam Stoelinga <sammiestoel@gmail.com>
Date: Tue, 21 Jan 2025 17:12:34 -0800
Subject: [PATCH] remove original gpu dockerfile contents

---
 Dockerfile | 13 ++-----------
 1 file changed, 2 insertions(+), 11 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index ec2d95f4..fac13013 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -97,17 +97,8 @@ COPY . .
 # GPU container spec.                                                          #
 ################################################################################
 
-# This causes INTERNAL: No valid engine configs for Matmul error
-# FROM base AS gpu
-#
-# RUN apt-get update && apt-get install -y ibverbs-utils
-# # TODO(markblee): Support extras.
-# ENV PIP_FIND_LINKS=https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-# RUN pip install .[core,gpu]
-# RUN pip install -U "jax[gpu]==0.4.37" "jax==0.4.37" "jaxlib==0.4.36" \
-#     -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-# COPY . .
-
+# Using `FROM base as GPU` causes INTERNAL: No valid engine configs for Matmul error.
+# So we're using the NVIDIA provided cuda image instead which works.
 FROM nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04 as gpu
 
 # Copy from original base