diff --git a/pylauncher.sh b/pylauncher.sh
index 21430bb8449db3702a9af3fcacb547014186b7cd..a81720ab1e1dbbff9d5e247faa4be9795eaf0936 100755
--- a/pylauncher.sh
+++ b/pylauncher.sh
@@ -10,7 +10,7 @@ shift 1
 
 source activate $environment
 
-if [[ $environment == "keras" || $environment == "pytorch" ]]; then
+if [[ $environment =~ ^keras.* || $environment =~ ^pytorch.* ]]; then
     export CUDA_HOME=/usr/local/cuda-8.0
     export CUDA_ROOT=$CUDA_HOME
 	export PATH=$CUDA_HOME/bin:$PATH