diff --git a/backends/vulkan/runtime/VulkanBackend.cpp b/backends/vulkan/runtime/VulkanBackend.cpp
index 51cb16764a..2621dc69d3 100644
--- a/backends/vulkan/runtime/VulkanBackend.cpp
+++ b/backends/vulkan/runtime/VulkanBackend.cpp
@@ -334,9 +334,6 @@ class GraphBuilder {
     }
 
     // Parse the operators
-    uint32_t last_prepack_node_ct = 0;
-    uint32_t last_execute_node_ct = 0;
-
     for (OpCallPtr op_call : *(flatbuffer_->chain())) {
       std::string op_name = op_call->name()->str();
       ET_CHECK_MSG(VK_HAS_OP(op_name), "Missing operator: %s", op_name.c_str());
@@ -351,22 +348,6 @@ class GraphBuilder {
 
       auto vkFn = VK_GET_OP_FN(op_name);
       vkFn(*compute_graph_, args);
-      if (compute_graph_->graphconfig().enable_querypool) {
-        for (uint32_t idx_prepack = last_prepack_node_ct;
-             idx_prepack < compute_graph_->prepack_nodes().size();
-             idx_prepack++) {
-          compute_graph_->prepack_nodes()[idx_prepack]->set_node_id(
-              op_call->node_id());
-        }
-        for (uint32_t idx_execute = last_execute_node_ct;
-             idx_execute < compute_graph_->execute_nodes().size();
-             idx_execute++) {
-          compute_graph_->execute_nodes()[idx_execute]->set_node_id(
-              op_call->node_id());
-        }
-        last_prepack_node_ct = compute_graph_->prepack_nodes().size();
-        last_execute_node_ct = compute_graph_->execute_nodes().size();
-      }
     }
 
     // Parse the outputs, which will be mostly tensors.  For some reason,
@@ -379,6 +360,15 @@ class GraphBuilder {
         compute_graph_->set_output_tensor(ref);
       }
     }
+
+    if (compute_graph_->graphconfig().enable_querypool) {
+      for (uint32_t i = 0; i < compute_graph_->prepack_nodes().size(); ++i) {
+        compute_graph_->prepack_nodes()[i]->set_node_id(i);
+      }
+      for (uint32_t i = 0; i < compute_graph_->execute_nodes().size(); ++i) {
+        compute_graph_->execute_nodes()[i]->set_node_id(i);
+      }
+    }
   }
 };