diff --git a/kernel/ethosu_network_info.c b/kernel/ethosu_network_info.c
index 5469b6c..5e7a1b9 100644
--- a/kernel/ethosu_network_info.c
+++ b/kernel/ethosu_network_info.c
@@ -115,19 +115,19 @@
 }
 
 void ethosu_network_info_rsp(struct ethosu_device *edev,
-			     struct ethosu_core_network_info_rsp *rsp)
+			     int msg_id,
+			     struct ethosu_core_msg_network_info_rsp *rsp)
 {
 	int ret;
-	int id = (int)rsp->user_arg;
 	struct ethosu_mailbox_msg *msg;
 	struct ethosu_network_info *info;
 	uint32_t i;
 
-	msg = ethosu_mailbox_find(&edev->mailbox, id);
+	msg = ethosu_mailbox_find(&edev->mailbox, msg_id);
 	if (IS_ERR(msg)) {
 		dev_warn(edev->dev,
 			 "Id for network info msg not found. msg.id=0x%x\n",
-			 id);
+			 msg_id);
 
 		return;
 	}