From 1a526cc34a1c104d98ab60e090e54252c6fedfc3 Mon Sep 17 00:00:00 2001
From: Yash Shah <55116947+yashs97@users.noreply.github.com>
Date: Thu, 26 Oct 2023 17:16:42 -0700
Subject: [PATCH] remove tbp nightly from setup.sh (#226)

---
 requirements.txt | 1 +
 setup.sh         | 6 ------
 2 files changed, 1 insertion(+), 6 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index ed5c9a1d9..704122695 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -17,3 +17,4 @@ tensorflow==2.13.*
 tensorflow-datasets
 tensorflow-text
 tensorboardx
+tensorboard-plugin-profile
diff --git a/setup.sh b/setup.sh
index 3a18c4fa9..deacb3f59 100644
--- a/setup.sh
+++ b/setup.sh
@@ -104,8 +104,6 @@ if [[ "$MODE" == "stable" || ! -v MODE ]]; then
         # Copy libtpu.so from GCS path
         gsutil cp "$LIBTPU_GCS_PATH" "$libtpu_path"
     fi
-    echo "Installing stable tensorboard plugin profile"
-    pip3 install tensorboard-plugin-profile --upgrade
 elif [[ $MODE == "nightly" ]]; then 
 # Nightly mode
     echo "Installing jax-head, jaxlib-nightly"
@@ -126,8 +124,6 @@ elif [[ $MODE == "nightly" ]]; then
         echo "Installing libtpu-nightly"
         pip3 install libtpu-nightly -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -U --pre
     fi
-    echo "Installing nightly tensorboard plugin profile"
-    pip3 install tbp-nightly --upgrade
 elif [[ $MODE == "head" ]]; then 
 # Head mode
     if [[ -n "$LIBTPU_GCS_PATH" ]]; then
@@ -154,8 +150,6 @@ elif [[ $MODE == "head" ]]; then
     cd $HOME/jax
     python3 build/build.py --enable_tpu --bazel_options="--override_repository=xla=$HOME/xla"
     pip3 install dist/jaxlib-*-cp*-manylinux2014_x86_64.whl --force-reinstall --no-deps
-    echo "Installing nightly tensorboard plugin profile"
-    pip3 install tbp-nightly --upgrade
 else
     echo -e "\n\nError: You can only set MODE to [stable,nightly,head,libtpu-only].\n\n"
     exit 1