diff --git a/README.md b/README.md
index 401dca31..2826cb0f 100644
--- a/README.md
+++ b/README.md
@@ -6,8 +6,7 @@ Clone the repository
 git clone git@github.com:awni/ecg.git
 ```
 
-
-Install `virtualenv` with
+If you don't have `virtualenv`, install it with
 
 ```
 pip install virtualenv
@@ -22,22 +21,14 @@ source ecg_env/bin/activate
 
 Install the requirements (this may take a few minutes).
 
-**NB** If you do not have a GPU, change line 67 in requirements.txt from 
-
+For CPU only support run
 ```
-tensorflow-gpu==1.0.1
+./setup.sh
 ```
 
-to
-
-```
-tensorflow==1.0.1
-```
-
-Now run
-
+To install with GPU support run
 ```
-pip install -r requirements.txt
+env TF=gpu ./setup.sh
 ```
 
 ## Training
@@ -63,7 +54,7 @@ python ecg/train.py
 After each epoch the model is saved in
 `ecg/saved/default/<experiment_id>/<model_id>.hdf5`.
 
-**NB** this model is only trained on 128 examples. This is far too few to see
+**NB**: this model is only trained on 128 examples. This is far too few to see
 good generalization performance, but the code should run and produce a valid
 model.
 
diff --git a/requirements.txt b/requirements.txt
index 17cbc405..58de79a9 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -64,7 +64,6 @@ six==1.10.0
 sklearn==0.0
 subprocess32==3.2.7
 tabulate==0.7.7
-tensorflow-gpu==1.0.1
 terminado==0.6
 testpath==0.3.1
 Theano==0.8.2
diff --git a/setup.sh b/setup.sh
new file mode 100755
index 00000000..7308b806
--- /dev/null
+++ b/setup.sh
@@ -0,0 +1,13 @@
+#!/bin/bash
+
+if [ -z "$TF" ]
+then
+    TF=tensorflow
+else
+    TF=tensorflow-gpu
+fi
+
+
+#pip install -r requirements.txt
+pip install --upgrade $TF==1.0.1 
+