diff --git a/src/strategies/avx2/dct-avx2.c b/src/strategies/avx2/dct-avx2.c
index ada4c9fe..61e08839 100644
--- a/src/strategies/avx2/dct-avx2.c
+++ b/src/strategies/avx2/dct-avx2.c
@@ -1634,7 +1634,7 @@ static void mts_idct_avx2(
 
   if (type_hor == DCT2 && type_ver == DCT2 && width == height)
   {
-    dct_func* idct_func = uvg_get_idct_func(width, color, tu->type);
+    dct_func* idct_func = uvg_get_idct_func(width, height, color, tu->type);
     idct_func(bitdepth, input, output);
   }
   else
diff --git a/src/strategies/generic/dct-generic.c b/src/strategies/generic/dct-generic.c
index c9074132..ff221180 100644
--- a/src/strategies/generic/dct-generic.c
+++ b/src/strategies/generic/dct-generic.c
@@ -2560,7 +2560,7 @@ static void mts_idct_generic(
 
   if (type_hor == DCT2 && type_ver == DCT2 && !tu->lfnst_idx && !tu->cr_lfnst_idx && width == height)
   {
-    dct_func *idct_func = uvg_get_idct_func(width, color, tu->type);
+    dct_func *idct_func = uvg_get_idct_func(width, height, color, tu->type);
     idct_func(bitdepth, input, output);
   }
   else
diff --git a/src/strategies/strategies-dct.c b/src/strategies/strategies-dct.c
index e7cc37e9..8441dfdd 100644
--- a/src/strategies/strategies-dct.c
+++ b/src/strategies/strategies-dct.c
@@ -99,7 +99,7 @@ dct_func * uvg_get_dct_func(int8_t width, int8_t height, color_t color, cu_type_
   if (width != height) {
     // Non-square block. Return generic dct for non-square blokcs.
     assert(false && "This should never be called at this point. Non-square stuff is done inside mts_dct function.");
-    return uvg_dct_non_square;
+    //return uvg_dct_non_square;
   }
   switch (width) {
   case 4:
@@ -128,8 +128,13 @@ dct_func * uvg_get_dct_func(int8_t width, int8_t height, color_t color, cu_type_
  *
  * \returns Pointer to the function.
  */
-dct_func * uvg_get_idct_func(int8_t width, color_t color, cu_type_t type)
+dct_func * uvg_get_idct_func(int8_t width, int8_t height, color_t color, cu_type_t type)
 {
+  if (width != height) {
+    // Non-square block. Return generic dct for non-square blokcs.
+    assert(false && "This should never be called at this point. Non-square stuff is done inside mts_idct function.");
+    //return uvg_idct_non_square;
+  }
   switch (width) {
   case 4:
     //if (color == COLOR_Y && type == CU_INTRA) {
diff --git a/src/strategies/strategies-dct.h b/src/strategies/strategies-dct.h
index 59e05084..b883b3e5 100644
--- a/src/strategies/strategies-dct.h
+++ b/src/strategies/strategies-dct.h
@@ -86,7 +86,7 @@ extern mts_idct_func* uvg_mts_idct;
 
 int uvg_strategy_register_dct(void* opaque, uint8_t bitdepth);
 dct_func * uvg_get_dct_func(int8_t width, int8_t height, color_t color, cu_type_t type);
-dct_func * uvg_get_idct_func(int8_t width, color_t color, cu_type_t type);
+dct_func * uvg_get_idct_func(int8_t width, int8_t height, color_t color, cu_type_t type);
 
 
 
diff --git a/src/transform.c b/src/transform.c
index 542ab3de..4a1b4042 100644
--- a/src/transform.c
+++ b/src/transform.c
@@ -271,14 +271,13 @@ void uvg_itransform2d(const encoder_control_t * const encoder,
                       color_t color,
                       const cu_info_t *tu)
 {
-  if (encoder->cfg.mts)
+  if (encoder->cfg.mts || block_width != block_height)
   {
     uvg_mts_idct(encoder->bitdepth, color, tu, block_width, block_height, coeff, block, encoder->cfg.mts);
   }
   else
   {
-    // ISP_TODO: block height
-    dct_func *idct_func = uvg_get_idct_func(block_width, color, tu->type);
+    dct_func *idct_func = uvg_get_idct_func(block_width, block_height, color, tu->type);
     idct_func(encoder->bitdepth, coeff, block);
   }
 }