diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 0f2eb3629070..59c4f2963417 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -3922,8 +3922,10 @@ pick_next_task(struct rq *rq, struct task_struct *prev, struct rq_flags *rf)
 			goto restart;
 
 		/* Assumes fair_sched_class->next == idle_sched_class */
-		if (unlikely(!p))
-			p = idle_sched_class.pick_next_task(rq, prev, rf);
+		if (unlikely(!p)) {
+			put_prev_task(rq, prev);
+			p = idle_sched_class.pick_next_task(rq, NULL, NULL);
+		}
 
 		return p;
 	}
diff --git a/kernel/sched/idle.c b/kernel/sched/idle.c
index f65ef1e2f204..179d1d4ac5a6 100644
--- a/kernel/sched/idle.c
+++ b/kernel/sched/idle.c
@@ -396,8 +396,7 @@ pick_next_task_idle(struct rq *rq, struct task_struct *prev, struct rq_flags *rf
 {
 	struct task_struct *next = rq->idle;
 
-	if (prev)
-		put_prev_task(rq, prev);
+	WARN_ON_ONCE(prev || rf);
 
 	set_next_task_idle(rq, next);