Update update_torch_wheel script

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

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"