from tensorflow.contrib.eager.python.examples.l2hmc import l2hmc

      learning_rate=.0003,
      n_warmup_iters=3)


def step(dynamics, optimizer, samples):
  loss, grads, samples, _ = l2hmc.loss_and_grads(
      dynamics, samples, loss_fn=l2hmc.compute_loss)
  optimizer.apply_gradients(zip(grads, dynamics.variables))

  return loss, samples
