import tensorflow as tf

          optimizer,
          n_iters=hparams.n_warmup_iters,
          n_samples=hparams.n_samples,
          step_fn=step_fn)

      # Training
      samples = tf.random_normal(
          shape=[hparams.n_samples, hparams.x_dim], dtype=tf.float32)
      start_time = time.time()
      fit(dynamics,
          samples,
          optimizer,
