diff --git a/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage-steps.yml b/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage-steps.yml
index 9b65ddbfdf3df..fc163d17e44a9 100644
--- a/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage-steps.yml
+++ b/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage-steps.yml
@@ -165,7 +165,7 @@ stages:
             set -e -x
             whlfilename=$(ls $(Build.ArtifactStagingDirectory)/Release/dist/*.whl | head -n 1)  ; \
             echo $whlfilename ; du -sh $whlfilename ; \
-            (( $(wc -c < "$whlfilename") -  300*1024*1024 < 0 )) ||  ( echo 'Wheel size bigger than 300M'; exit 1)
+            (( $(wc -c < "$whlfilename") -  400*1024*1024 < 0 )) ||  ( echo 'Wheel size bigger than 400M'; exit 1)
           displayName: 'Check wheel size'
           continueOnError: true