diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index 6778847..21725d7 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -4886,17 +4886,18 @@
       int count16x16_16x16p = 0, count16x16_lp = 0;
       int count32x32 = 0;
       for (i = 0; i < TX_SIZE_CONTEXTS; ++i) {
-        count4x4 += counts->tx_size[0][i][TX_4X4];
-        count4x4 += counts->tx_size[1][i][TX_4X4];
-        count4x4 += counts->tx_size[2][i][TX_4X4];
+        // counts->tx_size[max_depth][context_idx][this_depth_level]
+        count4x4 += counts->tx_size[0][i][0];
+        count4x4 += counts->tx_size[1][i][0];
+        count4x4 += counts->tx_size[2][i][0];
 
-        count8x8_lp += counts->tx_size[1][i][TX_8X8];
-        count8x8_lp += counts->tx_size[2][i][TX_8X8];
-        count8x8_8x8p += counts->tx_size[0][i][TX_8X8];
+        count8x8_lp += counts->tx_size[1][i][1];
+        count8x8_lp += counts->tx_size[2][i][1];
+        count8x8_8x8p += counts->tx_size[0][i][1];
 
-        count16x16_16x16p += counts->tx_size[1][i][TX_16X16];
-        count16x16_lp += counts->tx_size[2][i][TX_16X16];
-        count32x32 += counts->tx_size[2][i][TX_32X32];
+        count16x16_16x16p += counts->tx_size[1][i][2];
+        count16x16_lp += counts->tx_size[2][i][2];
+        count32x32 += counts->tx_size[2][i][3];
       }
 #if CONFIG_EXT_TX && CONFIG_RECT_TX
       count4x4 += counts->tx_size_implied[0][TX_4X4];