diff --git a/demos/pytorch/distributed/build_pytorch_occlum_instance.sh b/demos/pytorch/distributed/build_pytorch_occlum_instance.sh
index b78427b2..b34981c8 100755
--- a/demos/pytorch/distributed/build_pytorch_occlum_instance.sh
+++ b/demos/pytorch/distributed/build_pytorch_occlum_instance.sh
@@ -34,7 +34,7 @@ function build_instance()
     fi
 
     new_json="$(jq '.resource_limits.user_space_size = "4000MB" |
-                    .resource_limits.kernel_space_heap_size = "256MB" |
+                    .resource_limits.kernel_space_heap_size = "400MB" |
                     .resource_limits.max_num_of_threads = 64 |
                     .env.untrusted += [ "MASTER_ADDR", "MASTER_PORT", "WORLD_SIZE", "RANK", "OMP_NUM_THREADS", "HOME" ] |
                     .env.default += ["GLOO_DEVICE_TRANSPORT=TCP_TLS"] |