diff --git a/tools/ci_build/github/linux/run_dockerbuild.sh b/tools/ci_build/github/linux/run_dockerbuild.sh
index f792825aba..4f0f10cf0c 100755
--- a/tools/ci_build/github/linux/run_dockerbuild.sh
+++ b/tools/ci_build/github/linux/run_dockerbuild.sh
@@ -107,6 +107,8 @@ fi
 
 if [ $BUILD_DEVICE = "cpu" ] || [ $BUILD_DEVICE = "ngraph" ] || [ $BUILD_DEVICE = "openvino" ] || [ $BUILD_DEVICE = "nnapi" ] || [ $BUILD_DEVICE = "arm" ]; then
     RUNTIME=
+elif [[ $BUILD_EXTR_PAR = *--enable_training_python_frontend_e2e_tests* ]]; then
+     RUNTIME="--gpus all --shm-size=256m"
 else
     RUNTIME="--gpus all"
 fi