From 8b509b25a65d9ae06a87f9d9b1ac5de3210504bf Mon Sep 17 00:00:00 2001
From: Daniel Sohn <jysohn@google.com>
Date: Thu, 20 Jun 2019 17:32:11 -0700
Subject: [PATCH] Update update_torch_wheel script

---
 README.md                              |  7 ++++++
 scripts/update_nightly_torch_wheels.sh | 19 +++++++++++++++
 scripts/update_torch_wheels.sh         | 32 --------------------------
 3 files changed, 26 insertions(+), 32 deletions(-)
 create mode 100755 scripts/update_nightly_torch_wheels.sh
 delete mode 100755 scripts/update_torch_wheels.sh

diff --git a/README.md b/README.md
index b09810330465..513b20df08bf 100644
--- a/README.md
+++ b/README.md
@@ -73,6 +73,13 @@ environment are distributed under: `gcr.io/tpu-pytorch/xla`.
   (pytorch-0.1)$ python test/test_train_mnist.py
   ```
 
+  To update the wheels `torch` and `torch_xla` to the latest nightly
+  distribution (only updates your pytorch-nightly conda env), run:
+  ```Shell
+  (vm)$ cd /usr/share/torch-xla-nightly/pytorch/xla
+  (vm)$ ./scripts/update_nightly_torch_wheels.sh
+  ```
+
 ---
 
 # How To Build And Run PyTorch For TPU
diff --git a/scripts/update_nightly_torch_wheels.sh b/scripts/update_nightly_torch_wheels.sh
new file mode 100755
index 000000000000..009eb3fac2ec
--- /dev/null
+++ b/scripts/update_nightly_torch_wheels.sh
@@ -0,0 +1,19 @@
+#!/bin/bash
+set -e
+set -x
+
+DIST_BUCKET="gs://tpu-pytorch/wheels"
+TORCH_WHEEL="torch-nightly-cp36-cp36m-linux_x86_64.whl"
+TORCH_XLA_WHEEL="torch_xla-nightly-cp36-cp36m-linux_x86_64.whl"
+
+# Activate pytorch-nightly conda env if not already in it
+if [ "$CONDA_DEFAULT_ENV" != "pytorch-nightly" ]; then
+  source activate pytorch-nightly
+fi
+
+gsutil cp "$DIST_BUCKET/$TORCH_WHEEL" /tmp/
+gsutil cp "$DIST_BUCKET/$TORCH_XLA_WHEEL" /tmp/
+pip install "/tmp/$TORCH_WHEEL"
+pip install "/tmp/$TORCH_XLA_WHEEL"
+
+rm -f "/tmp/$TORCH_WHEEL" "/tmp/$TORCH_XLA_WHEEL"
diff --git a/scripts/update_torch_wheels.sh b/scripts/update_torch_wheels.sh
deleted file mode 100755
index 0fe596972123..000000000000
--- a/scripts/update_torch_wheels.sh
+++ /dev/null
@@ -1,32 +0,0 @@
-#!/bin/bash
-# Must be YYYYMMDD format
-set -e
-DATE=$1
-
-DIST_BUCKET="gs://pytorch-tpu-releases/wheels"
-TORCH_WHEEL="torch-1.1.0a0+stable-cp35-cp35m-linux_x86_64.whl"
-TORCH_XLA_WHEEL="torch_xla-0.1+nightly-cp35-cp35m-linux_x86_64.whl"
-STABLE_TORCH_WHEEL="torch-20190513-1.1.0a0+358fb51-cp35-cp35m-linux_x86_64.whl"
-STABLE_TORCH_XLA_WHEEL="torch_xla-20190513-0.1+6c24e5b-cp35-cp35m-linux_x86_64.whl"
-
-while getopts 's' OPTION
-do
-  case $OPTION in
-    s)
-      echo 'Installing stable build'
-      TORCH_WHEEL=$STABLE_TORCH_WHEEL
-      TORCH_XLA_WHEEL=$STABLE_TORCH_XLA_WHEEL
-      ;;
-    *)
-      echo 'Usage: ./update_torch_wheel -s (optional)'
-      exit 1
-      ;;
-  esac
-done
-
-gsutil cp "$DIST_BUCKET/$TORCH_WHEEL" /tmp/
-gsutil cp "$DIST_BUCKET/$TORCH_XLA_WHEEL" /tmp/
-pip3 install "/tmp/$TORCH_WHEEL"
-pip3 install "/tmp/$TORCH_XLA_WHEEL"
-
-rm -f "/tmp/$TORCH_WHEEL" "/tmp/$TORCH_XLA_WHEEL"