From c17f05400871d9ef8a59a809bcd3c01349ff7e73 Mon Sep 17 00:00:00 2001
From: kee hyun an <keehyuna@nvidia.com>
Date: Fri, 6 Dec 2024 11:29:32 +0900
Subject: [PATCH] chore: changed Mutex scope

---
 core/runtime/execute_engine.cpp | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/core/runtime/execute_engine.cpp b/core/runtime/execute_engine.cpp
index 4c2dc5c376..280c805295 100644
--- a/core/runtime/execute_engine.cpp
+++ b/core/runtime/execute_engine.cpp
@@ -98,7 +98,9 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
   LOG_DEBUG(
       "Attempting to run engine (ID: " << compiled_engine->name
                                        << "); Hardware Compatible: " << compiled_engine->hardware_compatible);
-
+  // nvinfer1::IExecutionContext::enqueue is not thread safe and we need a mutex for it.
+  // Other IExecutionContext methods and runtime states should be in same scope as well
+  std::unique_lock<std::mutex> lock(compiled_engine->mu);
   if (compiled_engine->profile_execution) {
     std::stringstream ss;
     ss << "Execution profiling is enabled, find results here:" << std::endl;
@@ -176,10 +178,6 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
     }
   }
 
-  // nvinfer1::IExecutionContext::enqueue is not thread safe and we need a mutex for it.
-  // setTensorAddress should be in same scope to prevent inconsistent state.
-  std::unique_lock<std::mutex> lock(compiled_engine->mu);
-
   { // Input Setup
     std::unique_ptr<torch::autograd::profiler::RecordProfile> input_profiler_guard;
     if (compiled_engine->profile_execution) {