From 2d596aa1403a958a52eb470a2647486eb6803966 Mon Sep 17 00:00:00 2001
From: erogol <erogol@hotmail.com>
Date: Tue, 14 Jul 2020 17:48:17 +0200
Subject: [PATCH] tf_lite tacotron2 bug fix

---
 tf/convert_tacotron2_tflite.py | 4 +---
 tf/utils/tflite.py             | 1 +
 vocoder/tf/layers/pqmf.py      | 2 +-
 3 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/tf/convert_tacotron2_tflite.py b/tf/convert_tacotron2_tflite.py
index e06cac2b..fc46cc79 100644
--- a/tf/convert_tacotron2_tflite.py
+++ b/tf/convert_tacotron2_tflite.py
@@ -34,6 +34,4 @@ model = load_checkpoint(model, args.tf_model)
 model.decoder.set_max_decoder_steps(1000)
 
 # create tflite model
-tflite_model = convert_tacotron2_to_tflite(model, output_path=args.output_path)
-
-print(f'Tflite Model size is {len(tflite_model) / (1024.0 * 1024.0)} MBs.')
+tflite_model = convert_tacotron2_to_tflite(model, output_path=args.output_path)
\ No newline at end of file
diff --git a/tf/utils/tflite.py b/tf/utils/tflite.py
index 6c37f170..5e684b30 100644
--- a/tf/utils/tflite.py
+++ b/tf/utils/tflite.py
@@ -16,6 +16,7 @@ def convert_tacotron2_to_tflite(model,
         tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS
     ]
     tflite_model = converter.convert()
+    print(f'Tflite Model size is {len(tflite_model) / (1024.0 * 1024.0)} MBs.')
     if output_path is not None:
         # same model binary if outputpath is provided
         with open(output_path, 'wb') as f:
diff --git a/vocoder/tf/layers/pqmf.py b/vocoder/tf/layers/pqmf.py
index 6c47dfc4..c018971f 100644
--- a/vocoder/tf/layers/pqmf.py
+++ b/vocoder/tf/layers/pqmf.py
@@ -51,7 +51,7 @@ class PQMF(tf.keras.layers.Layer):
 
     def synthesis(self, x):
         """
-        x : B x 1 x T
+        x : B x D x T
         """
         x = tf.transpose(x, perm=[0, 2, 1])
         x = tf.nn.conv1d_transpose(