diff --git a/src/encoder.c b/src/encoder.c
index 787d33cd..adc07c4c 100644
--- a/src/encoder.c
+++ b/src/encoder.c
@@ -2698,7 +2698,7 @@ void encode_coding_tree(encoder_state * const encoder_state,
     uint8_t intra_pred_mode[4] = {
       cur_cu->intra[0].mode, cur_cu->intra[1].mode,
       cur_cu->intra[2].mode, cur_cu->intra[3].mode };
-    uint8_t intra_pred_mode_chroma = 36; // 36 = Chroma derived from luma
+      uint8_t intra_pred_mode_chroma = cur_cu->intra[0].mode_chroma;
     int8_t intra_preds[4][3] = {{-1, -1, -1},{-1, -1, -1},{-1, -1, -1},{-1, -1, -1}};
     int8_t mpm_preds[4] = {-1, -1, -1, -1};
     int i, j;
@@ -2776,7 +2776,7 @@ void encode_coding_tree(encoder_state * const encoder_state,
       unsigned pred_mode = 5;
       unsigned chroma_pred_modes[4] = {0, 26, 10, 1};
 
-      if (intra_pred_mode_chroma == 36) {
+      if (intra_pred_mode_chroma == intra_pred_mode[0]) {
         pred_mode = 4;
       } else if (intra_pred_mode_chroma == 34) {
         // Angular 34 mode is possible only if intra pred mode is one of the
@@ -2791,6 +2791,10 @@ void encode_coding_tree(encoder_state * const encoder_state,
         }
       }
 
+      // pred_mode == 5 mean intra_pred_mode_chroma is something that can't
+      // be coded.
+      assert(pred_mode != 5);
+
       /**
        * Table 9-35 - Binarization for intra_chroma_pred_mode
        *   intra_chroma_pred_mode  bin_string
@@ -2992,11 +2996,6 @@ static void encode_transform_unit(encoder_state * const encoder_state,
       ctx_idx++;
       dir_mode = cur_cu->intra[0].mode_chroma;
 
-      if (dir_mode == 36) {
-        // TODO: support NxN
-        dir_mode = cur_cu->intra[0].mode;
-      }
-
       scan_idx = SCAN_DIAG;
 
       if (ctx_idx > 4 && ctx_idx < 7) { // if multiple scans supported for transform size
diff --git a/src/intra.c b/src/intra.c
index 376367e9..64297ae0 100644
--- a/src/intra.c
+++ b/src/intra.c
@@ -853,8 +853,6 @@ void intra_recon_lcu(encoder_state * const encoder_state, int x, int y, int dept
 
   int i = PU_INDEX(x >> 2, y >> 2);
 
-  cur_cu->intra[0].mode_chroma = 36; // TODO: Chroma intra prediction
-
   // Reconstruct chroma.
   if (!(x & 4 || y & 4)) {
     pixel *rec_shift_c  = &rec[width_c * 2 + 8 + 1];
@@ -866,7 +864,7 @@ void intra_recon_lcu(encoder_state * const encoder_state, int x, int y, int dept
                 width_c,
                 recbase_u,
                 rec_stride >> 1,
-                cur_cu->intra[0].mode_chroma != 36 ? cur_cu->intra[0].mode_chroma : cur_cu->intra[0].mode,
+                cur_cu->intra[0].mode_chroma,
                 1);
 
     intra_build_reference_border(encoder, x, y,(int16_t)width_c * 2 + 8, rec, (int16_t)width_c * 2 + 8, 2,
@@ -877,7 +875,7 @@ void intra_recon_lcu(encoder_state * const encoder_state, int x, int y, int dept
                 width_c,
                 recbase_v,
                 rec_stride >> 1,
-                cur_cu->intra[0].mode_chroma != 36 ? cur_cu->intra[0].mode_chroma : cur_cu->intra[0].mode,
+                cur_cu->intra[0].mode_chroma,
                 2);
   }
 
diff --git a/src/search.c b/src/search.c
index 65894b35..60f78bfa 100644
--- a/src/search.c
+++ b/src/search.c
@@ -718,6 +718,8 @@ static int search_cu_intra(encoder_state * const encoder_state,
     cur_cu->intra[pu_index].mode = (int8_t)mode;
     cur_cu->intra[pu_index].cost = cost;
     cur_cu->intra[pu_index].bitcost = bitcost;
+
+    cur_cu->intra[0].mode = cur_cu->intra[pu_index].mode;
   }
 
   return cur_cu->intra[PU_INDEX(x_px >> 2, y_px >> 2)].cost;
diff --git a/src/transform.c b/src/transform.c
index b55f6aea..1360917b 100644
--- a/src/transform.c
+++ b/src/transform.c
@@ -1146,9 +1146,6 @@ void quantize_lcu_chroma_residual(encoder_state * const encoder_state, int32_t x
     int chroma_depth = (depth == MAX_PU_DEPTH ? depth - 1 : depth);
     int chroma_width = LCU_WIDTH_C >> chroma_depth;
 
-    if (cur_cu->intra[0].mode_chroma == 36) {
-      cur_cu->intra[0].mode_chroma = cur_cu->intra[0].mode;
-    }
     scan_idx_chroma = get_scan_order(cur_cu->type, cur_cu->intra[0].mode_chroma, depth);
     if (quantize_residual(encoder_state, cur_cu, chroma_width, COLOR_U, scan_idx_chroma, tr_skip, LCU_WIDTH_C, LCU_WIDTH_C, base_u, recbase_u, recbase_u, orig_coeff_u)) {
       cbf_set(&cur_cu->cbf.u, depth);