diff --git a/kernel/ethosu_network_info.c b/kernel/ethosu_network_info.c
index 5e7a1b9..0e205db 100644
--- a/kernel/ethosu_network_info.c
+++ b/kernel/ethosu_network_info.c
@@ -31,10 +31,11 @@
 
 #define NETWORK_INFO_RESP_TIMEOUT_MS 3000
 
-static inline int ethosu_network_info_send(struct ethosu_network_info *info)
+static inline int ethosu_network_info_send(struct ethosu_network_info *info,
+					   struct ethosu_mailbox *mailbox)
 {
 	/* Send network info request to firmware */
-	return ethosu_mailbox_network_info_request(&info->edev->mailbox,
+	return ethosu_mailbox_network_info_request(mailbox,
 						   &info->msg,
 						   info->net->buf,
 						   info->net->index);
@@ -52,47 +53,49 @@
 	complete(&info->done);
 }
 
-int ethosu_network_info_request(struct ethosu_network *net,
+int ethosu_network_info_request(struct device *dev,
+				struct ethosu_mailbox *mailbox,
+				struct ethosu_network *net,
 				struct ethosu_uapi_network_info *uapi)
 {
 	struct ethosu_network_info *info;
 	int ret;
 	int timeout;
 
-	info = devm_kzalloc(net->edev->dev, sizeof(*info), GFP_KERNEL);
+	info = devm_kzalloc(dev, sizeof(*info), GFP_KERNEL);
 	if (!info)
 		return -ENOMEM;
 
-	info->edev = net->edev;
+	info->dev = dev;
 	info->net = net;
 	info->uapi = uapi;
 	init_completion(&info->done);
 	info->msg.fail = ethosu_network_info_fail;
 
-	ret = ethosu_mailbox_register(&info->edev->mailbox, &info->msg);
+	ret = ethosu_mailbox_register(mailbox, &info->msg);
 	if (ret < 0)
 		goto kfree;
 
 	/* Get reference to network */
 	ethosu_network_get(info->net);
 
-	ret = ethosu_network_info_send(info);
+	ret = ethosu_network_info_send(info, mailbox);
 	if (ret)
 		goto deregister;
 
-	dev_info(info->edev->dev,
+	dev_info(dev,
 		 "Network info create. info=0x%pK, net=0x%pK, msg.id=0x%x\n",
 		 info, info->net, info->msg.id);
 
 	/* Unlock the device mutex and wait for completion */
-	mutex_unlock(&info->edev->mutex);
+	device_unlock(dev);
 	timeout = wait_for_completion_timeout(&info->done,
 					      msecs_to_jiffies(
 						      NETWORK_INFO_RESP_TIMEOUT_MS));
-	mutex_lock(&info->edev->mutex);
+	device_lock(dev);
 
 	if (0 == timeout) {
-		dev_warn(info->edev->dev, "Network info timed out. info=0x%pK",
+		dev_warn(dev, "Network info timed out. info=0x%pK",
 			 info);
 
 		ret = -ETIME;
@@ -102,30 +105,31 @@
 	ret = info->errno;
 
 deregister:
-	ethosu_mailbox_deregister(&info->edev->mailbox, &info->msg);
+	ethosu_mailbox_deregister(mailbox, &info->msg);
 	ethosu_network_put(info->net);
 
 kfree:
-	dev_info(info->edev->dev,
+	dev_info(dev,
 		 "Network info destroy. info=0x%pK, msg.id=0x%x\n",
 		 info, info->msg.id);
-	devm_kfree(info->edev->dev, info);
+	devm_kfree(dev, info);
 
 	return ret;
 }
 
-void ethosu_network_info_rsp(struct ethosu_device *edev,
+void ethosu_network_info_rsp(struct ethosu_mailbox *mailbox,
 			     int msg_id,
 			     struct ethosu_core_msg_network_info_rsp *rsp)
 {
 	int ret;
+	struct device *dev = mailbox->dev;
 	struct ethosu_mailbox_msg *msg;
 	struct ethosu_network_info *info;
 	uint32_t i;
 
-	msg = ethosu_mailbox_find(&edev->mailbox, msg_id);
+	msg = ethosu_mailbox_find(mailbox, msg_id);
 	if (IS_ERR(msg)) {
-		dev_warn(edev->dev,
+		dev_warn(dev,
 			 "Id for network info msg not found. msg.id=0x%x\n",
 			 msg_id);