diff --git a/drivers/vhost/vhost.c b/drivers/vhost/vhost.c
index 22f775e514de..9963d5b80507 100644
--- a/drivers/vhost/vhost.c
+++ b/drivers/vhost/vhost.c
@@ -709,7 +709,8 @@ long vhost_dev_set_owner(struct vhost_dev *dev)
 	}
 
 	/* No owner, become one */
-	dev->mm = get_task_mm(current);
+	dev->mm = current->mm;
+	mmgrab(current->mm);
 	worker = kthread_create(vhost_worker, dev, "vhost-%d", current->pid);
 	if (IS_ERR(worker)) {
 		err = PTR_ERR(worker);
@@ -744,7 +745,7 @@ long vhost_dev_set_owner(struct vhost_dev *dev)
 	dev->worker = NULL;
 err_worker:
 	if (dev->mm)
-		mmput(dev->mm);
+		mmdrop(dev->mm);
 	dev->mm = NULL;
 err_mm:
 	return err;
@@ -963,7 +964,7 @@ void vhost_dev_cleanup(struct vhost_dev *dev)
 #if VHOST_ARCH_CAN_ACCEL_UACCESS
 		mmu_notifier_unregister(&dev->mmu_notifier, dev->mm);
 #endif
-		mmput(dev->mm);
+		mmdrop(dev->mm);
 	}
 #if VHOST_ARCH_CAN_ACCEL_UACCESS
 	for (i = 0; i < dev->nvqs; i++)