diff --git a/applications/message_handler/message_handler.hpp b/applications/message_handler/message_handler.hpp
index 5a195f0..36768ee 100644
--- a/applications/message_handler/message_handler.hpp
+++ b/applications/message_handler/message_handler.hpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2021 Arm Limited. All rights reserved.
+ * Copyright (c) 2020-2022 Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: Apache-2.0
  *
@@ -24,6 +24,7 @@
 #include "semphr.h"
 
 #include "message_queue.hpp"
+#include <ethosu_driver.h>
 #include <inference_process.hpp>
 #include <mailbox.hpp>
 
@@ -62,9 +63,14 @@
 private:
     void runInference(EthosU::ethosu_core_inference_req &req, EthosU::ethosu_core_inference_rsp &rsp);
 
+    friend void ::ethosu_inference_begin(struct ethosu_driver *drv, void *userArg);
+    friend void ::ethosu_inference_end(struct ethosu_driver *drv, void *userArg);
+
     QueueHandle_t inferenceQueue;
     QueueHandle_t outputQueue;
     InferenceProcess::InferenceProcess inference;
+    EthosU::ethosu_core_inference_req *currentReq;
+    EthosU::ethosu_core_inference_rsp *currentRsp;
 };
 
 struct OutputMessage {