From 55cb348fee19fe36983ebe88504e8d419a15f626 Mon Sep 17 00:00:00 2001 From: jorgep31415 Date: Thu, 12 Dec 2024 10:32:09 -0800 Subject: [PATCH] [ET-VK] Fix VulkanBackend `args` output index Pull Request resolved: https://github.com/pytorch/executorch/pull/7297 The one in the `VK_CHECK_COND` expression was wrong. Use a variable for easier tracking. ghstack-source-id: 257860131 @exported-using-ghexport Differential Revision: [D67126361](https://our.internmc.facebook.com/intern/diff/D67126361/) --- backends/vulkan/runtime/VulkanBackend.cpp | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/backends/vulkan/runtime/VulkanBackend.cpp b/backends/vulkan/runtime/VulkanBackend.cpp index 5093d61d2f7..51cb16764a3 100644 --- a/backends/vulkan/runtime/VulkanBackend.cpp +++ b/backends/vulkan/runtime/VulkanBackend.cpp @@ -582,16 +582,17 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface { compute_graph->execute(); for (size_t i = 0; i < compute_graph->outputs().size(); i++) { + const size_t o = i + num_inputs; const ValueRef oref = compute_graph->outputs()[i].value; if (compute_graph->val_is_tensor(oref)) { - VK_CHECK_COND(args[i]->isTensor()); - maybe_resize_output(compute_graph, i, args[num_inputs + i]->toTensor()); + VK_CHECK_COND(args[o]->isTensor()); + maybe_resize_output(compute_graph, i, args[o]->toTensor()); // args holds inputs directly followed by outputs, so the i'th output - // for compute_graph corresponds to the (i + num_inputs)'th arg + // for compute_graph corresponds to the o'th arg compute_graph->copy_from_staging( compute_graph->outputs()[i].staging, - args[num_inputs + i]->toTensor().mutable_data_ptr(), - args[num_inputs + i]->toTensor().numel()); + args[o]->toTensor().mutable_data_ptr(), + args[o]->toTensor().numel()); } else { VK_THROW( "Could not handle output with type ",