Correctly count the tx size cost in txfm_yrd

Change-Id: I6b050f3c0d34b6a8765dc69b64a5a9807678e504
diff --git a/av1/encoder/rdopt.c b/av1/encoder/rdopt.c
index 93baebf..7ce55fb 100644
--- a/av1/encoder/rdopt.c
+++ b/av1/encoder/rdopt.c
@@ -2372,8 +2372,10 @@
   const int is_inter = is_inter_block(mbmi);
   const int tx_select =
       cm->tx_mode == TX_MODE_SELECT && block_signals_txsize(mbmi->sb_type);
-
-  const int r_tx_size = tx_size_cost(cm, x, bs, tx_size);
+  int ctx = txfm_partition_context(
+      xd->above_txfm_context, xd->left_txfm_context, mbmi->sb_type, tx_size);
+  const int r_tx_size = is_inter ? x->txfm_partition_cost[ctx][0]
+                                 : tx_size_cost(cm, x, bs, tx_size);
 
   assert(IMPLIES(is_rect_tx(tx_size), is_rect_tx_allowed_bsize(bs)));