minor fix

Author: JamesDeng42

recml/core/training/jax_trainer.py

@@ -564,7 +564,7 @@ def _train_n_steps(self, train_iter, train_step, state, start_step, num_steps, s
     warmup_steps = 3
     total_examples_in_loop = 0
     valid_steps_in_loop = 0
-
+    loop_start_time = time.time()
 
     for step in range(start_step, start_step + num_steps):
       with jax.profiler.StepTraceAnnotation("train", step_num=step):