Update README.md

README.md

@@ -376,13 +376,6 @@ pip install witwidget
 jupyter nbextension install --py --symlink --sys-prefix witwidget
 jupyter nbextension enable --py --sys-prefix witwidget
 ```
-Note that if you use TensorFlow with GPU support (tensorflow-gpu), then you
-should instead install the GPU-compatible version of witwidget:
-```sh
-pip install witwidget-gpu
-jupyter nbextension install --py --symlink --sys-prefix witwidget
-jupyter nbextension enable --py --sys-prefix witwidget
-```
 
 Then, use it as seen at the bottom of the
 [What_If_Tool_Notebook_Usage.ipynb notebook](./What_If_Tool_Notebook_Usage.ipynb).
@@ -393,19 +386,21 @@ containing:
 ```
 !pip install witwidget
 ```
-For TensorFlow GPU support, use the `witwidget-gpu` package instead of `witwidget`.
 
 Then, use it as seen at the bottom of the
 [What_If_Tool_Notebook_Usage.ipynb notebook](https://colab.research.google.com/github/pair-code/what-if-tool/blob/master/What_If_Tool_Notebook_Usage.ipynb).
 
 ### How do I enable it for use in a JupyterLab or Cloud AI Platform notebook?
-Install and enable WIT for JupyterLab by running a cell containing:
+WIT has been tested in JupyterLab versions 1.x and 2.x.
+
+Install and enable WIT for JupyterLab 2.x by running a cell containing:
 ```
 !pip install witwidget
 !jupyter labextension install wit-widget
-!jupyter labextension install @jupyter-widgets/jupyterlab-manager
+!jupyter labextension install @jupyter-widgets/jupyterlab-manager@2
 ```
-For TensorFlow GPU support, use the `witwidget-gpu` package instead of `witwidget`.
+Note that you need to specify the correct version of jupyterlab-manager for your JupyterLab version as per https://www.npmjs.com/package/@jupyter-widgets/jupyterlab-manager.
+
 Note that you may need to run `!sudo jupyter labextension ...` commands depending on your notebook setup.
 
 Use of WIT after installation is the same as with the other notebook installations.