import tensorflow.contrib.eager as tfe

          x_dim=hparams.x_dim,
          minus_loglikelihood_fn=energy_fn,
          n_steps=hparams.n_steps,
          eps=hparams.eps)
      optimizer = tf.train.AdamOptimizer(learning_rate=hparams.learning_rate)
      step_fn = tfe.defun(step) if defun else step

      # Warmup to reduce initialization effect when timing
      warmup(
          dynamics,
