diff --git a/mm/memory.c b/mm/memory.c
index 40b7531ee8badb288823b64715265d0f16be397d..9b8a01d941cbc77dc0b1674fd817b54bb4f281f4 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -305,6 +305,7 @@ int __tlb_remove_page(struct mmu_gather *tlb, struct page *page)
 	if (batch->nr == batch->max) {
 		if (!tlb_next_batch(tlb))
 			return 0;
+		batch = tlb->active;
 	}
 	VM_BUG_ON(batch->nr > batch->max);