diff --git a/docker/scripts/startup.sh b/docker/scripts/startup.sh
index 527e8306a..1dbba9565 100644
--- a/docker/scripts/startup.sh
+++ b/docker/scripts/startup.sh
@@ -156,7 +156,7 @@ fi
 
 if [[ ${TEST} == "true" ]]; then
   cd algorithmic-efficiency
-  COMMAND="python3 tests/test_traindiffs.py"
+  COMMAND="python tests/test_traindiffs.py"
   echo $COMMAND
   eval $COMMAND
   exit
@@ -209,7 +209,7 @@ TUNING_RULESET_FLAG="--tuning_ruleset=${TUNING_RULESET}"
 
 # Set run command prefix depending on framework
 if [[ "${FRAMEWORK}" == "jax" ]]; then
-    COMMAND_PREFIX="python3"
+    COMMAND_PREFIX="python"
 else 
     COMMAND_PREFIX="torchrun --redirects 1:0,2:0,3:0,4:0,5:0,6:0,7:0 --standalone --nnodes=1 --nproc_per_node=8"
 fi