fix: using jax.random.key_data only when the workload is jax

Author: init-22

submission_runner.py

@@ -213,7 +213,9 @@ def train_once(
 ) -> Tuple[spec.Timing, Dict[str, Any]]:
   _reset_cuda_mem()
   data_rng, opt_init_rng, model_init_rng, rng = prng.split(rng, 4)
-  data_rng = jax.random.key_data(data_rng)
+
+  if FLAGS.framework == 'jax':
+    data_rng = jax.random.key_data(data_rng)
   # Workload setup.
   logging.info('Initializing dataset.')
   if hasattr(workload, '_eval_num_workers'):
@@ -345,7 +347,9 @@ def train_once(
 
     data_select_rng, update_rng, prep_eval_rng, eval_rng = \
       prng.split(step_rng, 4)
-    eval_rng = jax.random.key_data(eval_rng)
+
+    if FLAGS.framework == 'jax':
+      eval_rng = jax.random.key_data(eval_rng)
 
     with profiler.profile('Data selection'):
       batch = data_selection(workload,