diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp index e543218236c..5b2b48369a9 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -1451,6 +1451,10 @@ Error Method::reset_execution() { InvalidState, "Cannot reset until EndOfMethod has been reached."); step_state_ = StepState{0, 0}; + const auto n_input = inputs_size(); + for (size_t i = 0; i < n_input; ++i) { + input_set_[i] = false; + } return Error::Ok; } diff --git a/runtime/executor/test/kernel_integration_test.cpp b/runtime/executor/test/kernel_integration_test.cpp index 14fcb1c5260..8ad2f3d8d7e 100644 --- a/runtime/executor/test/kernel_integration_test.cpp +++ b/runtime/executor/test/kernel_integration_test.cpp @@ -294,6 +294,12 @@ TEST_F(KernelIntegrationTest, KernelHookIsCalled) { EXPECT_EQ(err, Error::Ok); EXPECT_EQ(control_->call_count, 1); + // Set up inputs again. + auto inputs_cleanup = executorch::extension::prepare_input_tensors(*method_); + ASSERT_EQ(inputs_cleanup.error(), Error::Ok); + auto input_err = method_->set_input(executorch::runtime::EValue(1.0), 2); + ASSERT_EQ(input_err, Error::Ok); + // Calling it again bumps the count. err = method_->execute(); EXPECT_EQ(err, Error::Ok); @@ -336,6 +342,12 @@ TEST_F(KernelIntegrationTest, DefaultPlatformMemoryAllocator) { EXPECT_EQ(control_->call_count, 1); EXPECT_EQ(control_->total_allocated_size, 4); + // Set up inputs again. + auto inputs_cleanup = executorch::extension::prepare_input_tensors(*method_); + ASSERT_EQ(inputs_cleanup.error(), Error::Ok); + auto input_err = method_->set_input(executorch::runtime::EValue(1.0), 2); + ASSERT_EQ(input_err, Error::Ok); + control_->temp_memory_size = 8; // This is not a simulation. This actually allocates memory, using the // default platform memory allocator. @@ -371,6 +383,12 @@ TEST_F(KernelTempMemoryAllocatorIntegrationTest, UsingTempMemoryAllocator) { EXPECT_EQ(temp_allocator_->number_of_resets, 1); EXPECT_EQ(temp_allocator_->currently_allocated_size, 0); + // Set up inputs again. + auto inputs_cleanup = executorch::extension::prepare_input_tensors(*method_); + ASSERT_EQ(inputs_cleanup.error(), Error::Ok); + auto input_err = method_->set_input(executorch::runtime::EValue(1.0), 2); + ASSERT_EQ(input_err, Error::Ok); + control_->temp_memory_size = 8; err = method_->execute(); EXPECT_EQ(err, Error::Ok); diff --git a/runtime/executor/test/method_test.cpp b/runtime/executor/test/method_test.cpp index 60f4e096bac..affd75b48c5 100644 --- a/runtime/executor/test/method_test.cpp +++ b/runtime/executor/test/method_test.cpp @@ -114,6 +114,12 @@ TEST_F(MethodTest, MoveTest) { Error err = method->execute(); ASSERT_EQ(err, Error::Ok); + // Set dummy inputs again. + auto input_cleanup2 = prepare_input_tensors(*method); + ASSERT_EQ(input_cleanup2.error(), Error::Ok); + input_err = method->set_input(executorch::runtime::EValue(1.0), 2); + ASSERT_EQ(input_err, Error::Ok); + // Move into a new Method. Method new_method(std::move(method.get()));