diff --git a/README.md b/README.md
index 4e98eb2..6b4ceb9 100644
--- a/README.md
+++ b/README.md
@@ -59,5 +59,5 @@ Synthesize speech
 -----------------
 
 ```sh
-python3 -m vietTTS.synthesizer --use-nat --text="hôm qua em tới trường" --output=clip.wav
+python3 -m vietTTS.synthesizer --use-nat --lexicon-file train_data/lexicon.txt --text="hôm qua em tới trường" --output=clip.wav
 ```
\ No newline at end of file
diff --git a/vietTTS/waveRNN/config.py b/vietTTS/waveRNN/config.py
index e1c5f88..f7005fb 100644
--- a/vietTTS/waveRNN/config.py
+++ b/vietTTS/waveRNN/config.py
@@ -27,6 +27,6 @@ class FLAGS:
 
   # training config
   ckpt_dir = Path('assets/reinfo/waveRNN')
-  wav_data_dir = Path('/tmp/reinfo/raw/pp')
+  wav_data_dir = Path('train_data')
   sample_rate = 16_000
   variance_loss_scale = 0.1  # regularization term
diff --git a/vietTTS/waveRNN/trainer.py b/vietTTS/waveRNN/trainer.py
index 29f3b91..e224923 100644
--- a/vietTTS/waveRNN/trainer.py
+++ b/vietTTS/waveRNN/trainer.py
@@ -61,7 +61,7 @@ def train():
   params, aux = net.init(rng, muinputs, mel)
   optimizer = make_optim(FLAGS._training_schedule[0].learning_rate)
   optim_state = optimizer.init(params)
-  training_step = 0
+  training_step = -1
 
   @jax.jit
   def update(params, aux, optim_state, batch, learning_rate):
diff --git a/vietTTS/waveRNN/utils.py b/vietTTS/waveRNN/utils.py
index c175dd1..740f769 100644
--- a/vietTTS/waveRNN/utils.py
+++ b/vietTTS/waveRNN/utils.py
@@ -83,7 +83,7 @@ def loop(mel, prev_state):
   return (out, reg, pr)
 
 
-regenerate_from_signal = jax.jit(regenerate_from_signal_.apply)
+regenerate_from_signal = jax.jit(regenerate_from_signal_.apply, static_argnums=[4])
 
 
 def gen_test_sample(params, aux, rng, test_clip, step=0, sr=16000):