diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index d8e12d5..032ae73 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -379,7 +379,7 @@
 
   if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide) return;
 
-  if (depth == 2) {
+  if (depth == MAX_VARTX_DEPTH) {
     txfm_partition_update(xd->above_txfm_context + tx_col,
                           xd->left_txfm_context + tx_row, tx_size);
     return;