move opt_init, ... = optimizers.adam(...) into the original location

Author: paul0403

demonstrations/tutorial_eqnn_force_field_catalyst_compiled.py

@@ -203,8 +203,6 @@ def inference(loss_data, opt_state):
 )
 
 ### training ###
-opt_init, opt_update, get_params = optimizers.adam(1e-2)
-
 np.random.seed(42)
 weights = np.zeros((num_qubits, D, B))
 weights[0] = np.random.uniform(0, np.pi, 1)
@@ -220,7 +218,7 @@ def inference(loss_data, opt_state):
 epsilon = jax.lax.stop_gradient(epsilon)  # comment if we wish to train the SB weights as well.
 
 
-
+opt_init, opt_update, get_params = optimizers.adam(1e-2)
 net_params = {"params": {"weights": weights, "alphas": alphas, "epsilon": epsilon}}
 opt_state = opt_init(net_params)
 running_loss = []