diff --git a/.github/workflows/demo_test.yml b/.github/workflows/demo_test.yml
index 1001b45e..20e35dd5 100644
--- a/.github/workflows/demo_test.yml
+++ b/.github/workflows/demo_test.yml
@@ -346,7 +346,7 @@ jobs:
     # FIXME: PyTorch can't exit normally in SIM mode
     - name: Kill the container
       run: |
-        sleep 240;
+        sleep 360;
         cat "$GITHUB_WORKSPACE/log";
         if grep -q Done "$GITHUB_WORKSPACE/log"; then
           docker stop pytorch_test
diff --git a/demos/pytorch/run_pytorch_on_occlum.sh b/demos/pytorch/run_pytorch_on_occlum.sh
index f701209b..7d964704 100755
--- a/demos/pytorch/run_pytorch_on_occlum.sh
+++ b/demos/pytorch/run_pytorch_on_occlum.sh
@@ -17,7 +17,7 @@ fi
 
 new_json="$(jq '.resource_limits.user_space_size = "6000MB" |
                 .resource_limits.kernel_space_heap_size = "256MB" |
-                .process.default_mmap_size = "4000MB" |
+                .resource_limits.max_num_of_threads = 64 |
                 .env.default += ["PYTHONHOME=/opt/python-occlum"]' Occlum.json)" && \
 echo "${new_json}" > Occlum.json
 occlum build