From 70dbb16ad10915b71d5ea63bc5391b4689b7503f Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <shoumikhin@meta.com>
Date: Thu, 26 Sep 2024 13:33:04 -0700
Subject: [PATCH] Small improvements for module usage.

Summary: .

Reviewed By: kirklandsign

Differential Revision: D63473870
---
 extension/apple/Benchmark/Tests/Tests.mm | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/extension/apple/Benchmark/Tests/Tests.mm b/extension/apple/Benchmark/Tests/Tests.mm
index 74388b49d54..7e483a47c83 100644
--- a/extension/apple/Benchmark/Tests/Tests.mm
+++ b/extension/apple/Benchmark/Tests/Tests.mm
@@ -99,7 +99,7 @@ + (void)initialize {
         IMP testForwardImplementation = imp_implementationWithBlock(^(
             id _self) {
           auto __block module = std::make_unique<Module>(modelPath.UTF8String);
-          XCTAssertEqual(module->load_method("forward"), Error::Ok);
+          XCTAssertEqual(module->load_forward(), Error::Ok);
 
           const auto method_meta = module->method_meta("forward");
           XCTAssertEqual(method_meta.error(), Error::Ok);
@@ -109,8 +109,6 @@ + (void)initialize {
 
           std::vector<TensorPtr> __block tensors;
           tensors.reserve(num_inputs);
-          std::vector<EValue> __block inputs;
-          inputs.reserve(num_inputs);
 
           for (auto index = 0; index < num_inputs; ++index) {
             const auto input_tag = method_meta->input_tag(index);
@@ -124,7 +122,7 @@ + (void)initialize {
               const auto sizes = tensor_meta->sizes();
               tensors.emplace_back(ones({sizes.begin(), sizes.end()},
                                         tensor_meta->scalar_type()));
-              inputs.emplace_back(tensors.back());
+              module.set_input(tensors.back(), index);
             } break;
             default:
               XCTFail("Unsupported tag %i at input %d", *input_tag, index);