From 7f3571321d27fc187aed5808623a4e9b0c5cfbde Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <shoumikhin@meta.com>
Date: Mon, 30 Sep 2024 21:14:10 -0700
Subject: [PATCH] Fix load time.

Summary: Once loaded Module does nothing next time, we need to recreate the Module again to truly load the method again.

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

diff --git a/extension/apple/Benchmark/Tests/Tests.mm b/extension/apple/Benchmark/Tests/Tests.mm
index 8c399da5744..2c443b97e22 100644
--- a/extension/apple/Benchmark/Tests/Tests.mm
+++ b/extension/apple/Benchmark/Tests/Tests.mm
@@ -77,15 +77,15 @@ + (void)initialize {
                                        modelName,
                                        deviceInfo]);
         IMP testLoadImplementation = imp_implementationWithBlock(^(id _self) {
-          auto __block module = std::make_unique<Module>(modelPath.UTF8String);
           [_self measureWithMetrics:@[
             [XCTClockMetric new],
             [XCTMemoryMetric new],
           ]
                             options:XCTMeasureOptions.defaultOptions
                               block:^{
-                                XCTAssertEqual(module->load_method("forward"),
-                                               Error::Ok);
+                                XCTAssertEqual(
+                                    Module(modelPath.UTF8String).load_forward(),
+                                    Error::Ok);
                               }];
         });
         class_addMethod(
@@ -107,7 +107,7 @@ + (void)initialize {
           const auto num_inputs = method_meta->num_inputs();
           XCTAssertGreaterThan(num_inputs, 0);
 
-          std::vector<TensorPtr> __block tensors;
+          std::vector<TensorPtr> tensors;
           tensors.reserve(num_inputs);
 
           for (auto index = 0; index < num_inputs; ++index) {