code: def plot_val_duration(step: int, batch, params, aux, rng):
      fn = FLAGS.ckpt_dir / f"duration_{step}.png"
      predicted_dur, gt_dur = predict_duration(params, aux, rng, batch)
      L = batch.lengths[0]
      x = np.arange(0, L) * 3
      plt.plot(predicted_dur[0, :L])
      plt.plot(gt_dur[0, :L])
      plt.legend(["predicted", "gt"])
      plt.title("Phoneme durations")
      plt.savefig(fn)
      plt.close()
  
file_path: underthesea/underthesea/pipeline/say/viettts_/nat/duration_trainer.py
imports: from functools import partial
  from typing import Deque
  import haiku
  import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  import optax
  from tqdm.auto import tqdm
  from vietTTS.nat.config import DurationInput
  from .config import FLAGS
  from .data_loader import textgrid_data_loader
  from .model import DurationModel
  from .utils import load_latest_ckpt, print_flags, save_ckpt
label_desc: Use image node_id 5 for calling a global function with example usage: plot_val_duration(step, batch, params, aux, rng) without return types