2 files changed, 3 insertions, 0 deletions

diff --git a/kernel/fork.c b/kernel/fork.c
index f71b67dc156d..242c8c93d285 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -595,6 +595,8 @@ static void check_mm(struct mm_struct *mm)
 void __mmdrop(struct mm_struct *mm)
 {
         BUG_ON(mm == &init_mm);
+        WARN_ON_ONCE(mm == current->mm);
+        WARN_ON_ONCE(mm == current->active_mm);
         mm_free_pgd(mm);
         destroy_context(mm);
         hmm_mm_destroy(mm);
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 28b68995a417..e8afd6086f23 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -5560,6 +5560,7 @@ void idle_task_exit(void)
 
         if (mm != &init_mm) {
                 switch_mm(mm, &init_mm, current);
+                current->active_mm = &init_mm;
                 finish_arch_post_lock_switch();
         }
         mmdrop(mm);