diff --git a/kernel/ethosu_mailbox.h b/kernel/ethosu_mailbox.h
index c5865d0..5a77b54 100644
--- a/kernel/ethosu_mailbox.h
+++ b/kernel/ethosu_mailbox.h
@@ -100,8 +100,10 @@
  */
 int ethosu_mailbox_inference(struct ethosu_mailbox *mbox,
 			     void *user_arg,
-			     struct ethosu_buffer *ifm,
-			     struct ethosu_buffer *ofm,
+			     uint32_t ifm_count,
+			     struct ethosu_buffer **ifm,
+			     uint32_t ofm_count,
+			     struct ethosu_buffer **ofm,
 			     struct ethosu_buffer *network);
 
 #endif /* ETHOSU_MAILBOX_H */