diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index 12e8a81..b7b8208 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -4666,17 +4666,17 @@
   }
 
 #if CONFIG_EXT_PARTITION_TYPES
-  prob_diff_update(av1_partition_tree, fc->partition_prob[0],
-                   counts->partition[0], PARTITION_TYPES, probwt, header_bc);
-  for (i = 1; i < PARTITION_CONTEXTS_PRIMARY; ++i)
+  for (i = 0; i < PARTITION_PLOFFSET; ++i)
+    prob_diff_update(av1_partition_tree, fc->partition_prob[i],
+                     counts->partition[i], PARTITION_TYPES, probwt, header_bc);
+  for (; i < PARTITION_CONTEXTS_PRIMARY; ++i)
     prob_diff_update(av1_ext_partition_tree, fc->partition_prob[i],
                      counts->partition[i], EXT_PARTITION_TYPES, probwt,
                      header_bc);
 #else
-  for (i = 0; i < PARTITION_CONTEXTS_PRIMARY; ++i) {
+  for (i = 0; i < PARTITION_CONTEXTS_PRIMARY; ++i)
     prob_diff_update(av1_partition_tree, fc->partition_prob[i],
                      counts->partition[i], PARTITION_TYPES, probwt, header_bc);
-  }
 #endif  // CONFIG_EXT_PARTITION_TYPES
 #if CONFIG_UNPOISON_PARTITION_CTX
   for (; i < PARTITION_CONTEXTS_PRIMARY + PARTITION_BLOCK_SIZES; ++i) {
diff --git a/av1/encoder/rd.c b/av1/encoder/rd.c
index 76d18ea..d6561fc 100644
--- a/av1/encoder/rd.c
+++ b/av1/encoder/rd.c
@@ -405,9 +405,10 @@
       cpi->partition_cost[0][PARTITION_SPLIT] = 0;
 #endif
 #if CONFIG_EXT_PARTITION_TYPES
-      av1_cost_tokens(cpi->partition_cost[CONFIG_UNPOISON_PARTITION_CTX],
-                      cm->fc->partition_prob[0], av1_partition_tree);
-      for (i = 1; i < PARTITION_CONTEXTS_PRIMARY; ++i)
+      for (i = 0; i < PARTITION_PLOFFSET; ++i)
+        av1_cost_tokens(cpi->partition_cost[CONFIG_UNPOISON_PARTITION_CTX + i],
+                        cm->fc->partition_prob[i], av1_partition_tree);
+      for (; i < PARTITION_CONTEXTS_PRIMARY; ++i)
         av1_cost_tokens(cpi->partition_cost[CONFIG_UNPOISON_PARTITION_CTX + i],
                         cm->fc->partition_prob[i], av1_ext_partition_tree);
 #else