diff --git a/arch/powerpc/kernel/mce_power.c b/arch/powerpc/kernel/mce_power.c
index de242b4bbd20b5c019c1d19072d3ee30818b9c26..f913139bb0c2c9cdd67a20514a453e0b223c314b 100644
--- a/arch/powerpc/kernel/mce_power.c
+++ b/arch/powerpc/kernel/mce_power.c
@@ -72,10 +72,14 @@ void __flush_tlb_power8(unsigned int action)
 
 void __flush_tlb_power9(unsigned int action)
 {
+	unsigned int num_sets;
+
 	if (radix_enabled())
-		flush_tlb_206(POWER9_TLB_SETS_RADIX, action);
+		num_sets = POWER9_TLB_SETS_RADIX;
+	else
+		num_sets = POWER9_TLB_SETS_HASH;
 
-	flush_tlb_206(POWER9_TLB_SETS_HASH, action);
+	flush_tlb_206(num_sets, action);
 }