Merge pull request #49 from deepnote/hannes/pla-3506-fix-gpu-images-to-work-with-toolkit

Fix GPU images to work with toolkit

Author: hc2p

gpu/setup.python.sh

@@ -1,8 +1,8 @@
 #!/bin/bash
 
-# 
+#
 # This script allows us to install a python version of choice
-# 
+#
 
 PYTHON_VER=$1
 TF_VERSION=$2
@@ -30,10 +30,7 @@ ln -s /usr/bin/python3 /usr/bin/python
 
 # We need to remove some pre-installed pip packages that end up clashing with our compute-deps
 pip uninstall -y cryptography
-pip install numpy==1.23.4
-
-# install tensorflow
-pip install tensorflow==$TF_VERSION
+pip install numpy==1.23.4 tensorflow==$TF_VERSION -c https://tk.deepnote.com/constraints${PYTHON_VER}.txt
 
 # create the virtual environment in the home directory in the Dockerfile
 # using our upgraded python version