diff --git a/pylauncher.sh b/pylauncher.sh
index ae96ee3149925c99b1d752c065cce445bcb09e8e..21430bb8449db3702a9af3fcacb547014186b7cd 100755
--- a/pylauncher.sh
+++ b/pylauncher.sh
@@ -7,11 +7,10 @@ fi
 
 environment="$1"
 shift 1
-# command="$2"
 
 source activate $environment
 
-if [[ $environment == "keras" ]]; 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