diff --git a/drivers/vhost/net.c b/drivers/vhost/net.c
index e34e195..a3645bd 100644
--- a/drivers/vhost/net.c
+++ b/drivers/vhost/net.c
@@ -967,14 +967,20 @@
 	struct socket *tx_sock = NULL;
 	struct socket *rx_sock = NULL;
 	long err;
+	struct vhost_memory *memory;
 
 	mutex_lock(&n->dev.mutex);
 	err = vhost_dev_check_owner(&n->dev);
 	if (err)
 		goto done;
+	memory = vhost_dev_reset_owner_prepare();
+	if (!memory) {
+		err = -ENOMEM;
+		goto done;
+	}
 	vhost_net_stop(n, &tx_sock, &rx_sock);
 	vhost_net_flush(n);
-	err = vhost_dev_reset_owner(&n->dev);
+	vhost_dev_reset_owner(&n->dev, memory);
 	vhost_net_vq_reset(n);
 done:
 	mutex_unlock(&n->dev.mutex);