diff --git a/drivers/vhost/vhost.c b/drivers/vhost/vhost.c
index dc9301d31f12..f7bbbc1b6dca 100644
--- a/drivers/vhost/vhost.c
+++ b/drivers/vhost/vhost.c
@@ -2065,8 +2065,10 @@ static long vhost_vring_set_num_addr(struct vhost_dev *d,
 	/* Unregister MMU notifer to allow invalidation callback
 	 * can access vq->uaddrs[] without holding a lock.
 	 */
-	if (d->mm)
+	if (d->mm) {
 		mmu_notifier_unregister(&d->mmu_notifier, d->mm);
+		mmu_notifier_register(&d->mmu_notifier, d->mm);
+	}
 
 	vhost_uninit_vq_maps(vq);
 #endif
@@ -2085,8 +2087,6 @@ static long vhost_vring_set_num_addr(struct vhost_dev *d,
 #if VHOST_ARCH_CAN_ACCEL_UACCESS
 	vhost_setup_vq_uaddr(vq);
 
-	if (d->mm)
-		mmu_notifier_register(&d->mmu_notifier, d->mm);
 #endif
 
 	mutex_unlock(&vq->mutex);