Add the k8s extra when installing JAX (#1254)

NVIDIA · Jan 17, 2025 · 9ba6461 · 9ba6461
1 parent eb6d0d2
commit 9ba6461
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/.github/container/Dockerfile.jax b/.github/container/Dockerfile.jax
@@ -102,7 +102,7 @@ RUN <<"EOF" bash -ex
 for component in $(ls ${BUILD_PATH_JAXLIB}); do
     echo "-e file://${BUILD_PATH_JAXLIB}/${component}" >> /opt/pip-tools.d/requirements-jax.in;
 done
-echo "-e file://${SRC_PATH_JAX}" >> /opt/pip-tools.d/requirements-jax.in
+echo "-e file://${SRC_PATH_JAX}[k8s]" >> /opt/pip-tools.d/requirements-jax.in
 echo "numpy<2.0.0" >> /opt/pip-tools.d/requirements-jax.in
 EOF