diff --git a/av1/decoder/decodeframe.c b/av1/decoder/decodeframe.c
index d81a21a..12ecc44 100644
--- a/av1/decoder/decodeframe.c
+++ b/av1/decoder/decodeframe.c
@@ -3085,6 +3085,15 @@
   pool->frame_bufs[cm->new_fb_idx].buf.render_height = cm->render_height;
 }
 
+static void setup_sb_size(AV1_COMMON *cm, struct aom_read_bit_buffer *rb) {
+  (void)rb;
+#if CONFIG_EXT_PARTITION
+  set_sb_size(cm, aom_rb_read_bit(rb) ? BLOCK_128X128 : BLOCK_64X64);
+#else
+  set_sb_size(cm, BLOCK_64X64);
+#endif  // CONFIG_EXT_PARTITION
+}
+
 static INLINE int valid_ref_frame_img_fmt(aom_bit_depth_t ref_bit_depth,
                                           int ref_xss, int ref_yss,
                                           aom_bit_depth_t this_bit_depth,
@@ -4720,6 +4729,8 @@
     }
 
     setup_frame_size(cm, rb);
+    setup_sb_size(cm, rb);
+
     if (pbi->need_resync) {
       memset(&cm->ref_frame_map, -1, sizeof(cm->ref_frame_map));
       pbi->need_resync = 0;
@@ -4967,12 +4978,6 @@
   if (frame_is_intra_only(cm) || cm->error_resilient_mode)
     av1_setup_past_independence(cm);
 
-#if CONFIG_EXT_PARTITION
-  set_sb_size(cm, aom_rb_read_bit(rb) ? BLOCK_128X128 : BLOCK_64X64);
-#else
-  set_sb_size(cm, BLOCK_64X64);
-#endif  // CONFIG_EXT_PARTITION
-
   setup_loopfilter(cm, rb);
   setup_quantization(cm, rb);
   xd->bd = (int)cm->bit_depth;
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index d5b9380..d28244b 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -4319,6 +4319,20 @@
 }
 #endif
 
+static void write_sb_size(const AV1_COMMON *cm,
+                          struct aom_write_bit_buffer *wb) {
+  (void)cm;
+  (void)wb;
+  assert(cm->mib_size == mi_size_wide[cm->sb_size]);
+  assert(cm->mib_size == 1 << cm->mib_size_log2);
+#if CONFIG_EXT_PARTITION
+  assert(cm->sb_size == BLOCK_128X128 || cm->sb_size == BLOCK_64X64);
+  aom_wb_write_bit(wb, cm->sb_size == BLOCK_128X128 ? 1 : 0);
+#else
+  assert(cm->sb_size == BLOCK_64X64);
+#endif  // CONFIG_EXT_PARTITION
+}
+
 #if CONFIG_EXT_INTER
 static void write_compound_tools(const AV1_COMMON *cm,
                                  struct aom_write_bit_buffer *wb) {
@@ -4538,6 +4552,8 @@
     write_sync_code(wb);
     write_bitdepth_colorspace_sampling(cm, wb);
     write_frame_size(cm, wb);
+    write_sb_size(cm, wb);
+
 #if CONFIG_ANS && ANS_MAX_SYMBOLS
     assert(cpi->common.ans_window_size_log2 >= 8);
     assert(cpi->common.ans_window_size_log2 < 24);
@@ -4684,15 +4700,6 @@
 #if !CONFIG_NO_FRAME_CONTEXT_SIGNALING
   aom_wb_write_literal(wb, cm->frame_context_idx, FRAME_CONTEXTS_LOG2);
 #endif
-  assert(cm->mib_size == mi_size_wide[cm->sb_size]);
-  assert(cm->mib_size == 1 << cm->mib_size_log2);
-#if CONFIG_EXT_PARTITION
-  assert(cm->sb_size == BLOCK_128X128 || cm->sb_size == BLOCK_64X64);
-  aom_wb_write_bit(wb, cm->sb_size == BLOCK_128X128 ? 1 : 0);
-#else
-  assert(cm->sb_size == BLOCK_64X64);
-#endif  // CONFIG_EXT_PARTITION
-
   encode_loopfilter(cm, wb);
   encode_quantization(cm, wb);
   encode_segmentation(cm, xd, wb);
@@ -4829,8 +4836,9 @@
 
   if (cm->frame_type == KEY_FRAME) {
     write_sync_code(wb);
-
     write_frame_size(cm, wb);
+    write_sb_size(cm, wb);
+
 #if CONFIG_ANS && ANS_MAX_SYMBOLS
     assert(cpi->common.ans_window_size_log2 >= 8);
     assert(cpi->common.ans_window_size_log2 < 24);
@@ -5041,15 +5049,6 @@
 #if !CONFIG_NO_FRAME_CONTEXT_SIGNALING
   aom_wb_write_literal(wb, cm->frame_context_idx, FRAME_CONTEXTS_LOG2);
 #endif
-  assert(cm->mib_size == mi_size_wide[cm->sb_size]);
-  assert(cm->mib_size == 1 << cm->mib_size_log2);
-#if CONFIG_EXT_PARTITION
-  assert(cm->sb_size == BLOCK_128X128 || cm->sb_size == BLOCK_64X64);
-  aom_wb_write_bit(wb, cm->sb_size == BLOCK_128X128 ? 1 : 0);
-#else
-  assert(cm->sb_size == BLOCK_64X64);
-#endif  // CONFIG_EXT_PARTITION
-
   encode_loopfilter(cm, wb);
   encode_quantization(cm, wb);
   encode_segmentation(cm, xd, wb);
diff --git a/av1/encoder/encoder.c b/av1/encoder/encoder.c
index a8d8989..fa4093e 100644
--- a/av1/encoder/encoder.c
+++ b/av1/encoder/encoder.c
@@ -344,6 +344,7 @@
     cpi->refresh_golden_frame = 1;
     cpi->refresh_alt_ref_frame = 1;
     av1_zero(cpi->interp_filter_selected);
+    set_sb_size(cm, select_sb_size(cpi));
   } else {
 #if CONFIG_NO_FRAME_CONTEXT_SIGNALING
     if (frame_is_intra_only(cm) || cm->error_resilient_mode ||
@@ -380,8 +381,6 @@
 #endif  // CONFIG_NO_FRAME_CONTEXT_SIGNALING
 
   cpi->vaq_refresh = 0;
-
-  set_sb_size(cm, select_sb_size(cpi));
 }
 
 static void enc_setup_mi(AV1_COMMON *cm) {