diff --git a/vp9/encoder/vp9_encodeframe.c b/vp9/encoder/vp9_encodeframe.c
index 7c23973b84c7e08e9272c920008ff313afecd3ce..1e5890cda07b6d7e60371104453422d74ae8079c 100644
--- a/vp9/encoder/vp9_encodeframe.c
+++ b/vp9/encoder/vp9_encodeframe.c
@@ -1405,7 +1405,7 @@ static int sb_has_motion(const VP9_COMMON *cm, MODE_INFO **prev_mi_8x8) {
   return 0;
 }
 
-static void update_state_rt(VP9_COMP *cpi, const PICK_MODE_CONTEXT *ctx,
+static void update_state_rt(VP9_COMP *cpi, PICK_MODE_CONTEXT *ctx,
                             int mi_row, int mi_col, int bsize) {
   VP9_COMMON *const cm = &cpi->common;
   MACROBLOCK *const x = &cpi->mb;
@@ -1413,15 +1413,11 @@ static void update_state_rt(VP9_COMP *cpi, const PICK_MODE_CONTEXT *ctx,
   MB_MODE_INFO *const mbmi = &xd->mi_8x8[0]->mbmi;
   const struct segmentation *const seg = &cm->seg;
 
-  // TODO(jingning) We might need PICK_MODE_CONTEXT to buffer coding modes
-  // associated with variable block sizes. Otherwise, remove this ctx
-  // from argument list.
-  (void)ctx;
-
   *(xd->mi_8x8[0]) = ctx->mic;
 
-  // Check for reseting segment_id and update cyclic map.
-  if (cpi->oxcf.aq_mode == CYCLIC_REFRESH_AQ && seg->enabled) {
+  // For in frame adaptive Q, check for reseting the segment_id and updating
+  // the cyclic refresh map.
+  if ((cpi->oxcf.aq_mode == CYCLIC_REFRESH_AQ) && seg->enabled) {
     vp9_cyclic_refresh_update_segment(cpi, &xd->mi_8x8[0]->mbmi,
                                       mi_row, mi_col, bsize, 1);
     vp9_init_plane_quantizers(cpi, x);
@@ -1448,6 +1444,7 @@ static void encode_b_rt(VP9_COMP *cpi, const TileInfo *const tile,
     if (x->ab_index > 0)
       return;
   }
+
   set_offsets(cpi, tile, mi_row, mi_col, bsize);
   update_state_rt(cpi, get_block_context(x, bsize), mi_row, mi_col, bsize);
 
diff --git a/vp9/encoder/vp9_pickmode.c b/vp9/encoder/vp9_pickmode.c
index bd3be907b3a8d98556a276834172d4280dca3e0b..c500a01631842f134807c7a3b1118c6cb5976f3a 100644
--- a/vp9/encoder/vp9_pickmode.c
+++ b/vp9/encoder/vp9_pickmode.c
@@ -255,7 +255,7 @@ int64_t vp9_pick_inter_mode(VP9_COMP *cpi, MACROBLOCK *x,
   mbmi->interp_filter = cpi->common.interp_filter == SWITCHABLE ?
                         EIGHTTAP : cpi->common.interp_filter;
   mbmi->skip = 0;
-  mbmi->segment_id = 0;
+  mbmi->segment_id = segment_id;
   xd->interp_kernel = vp9_get_interp_kernel(mbmi->interp_filter);
 
   for (ref_frame = LAST_FRAME; ref_frame <= LAST_FRAME ; ++ref_frame) {