revert changes to submission runner for prng key

Author: priyakasimbeg

submission_runner.py

@@ -214,8 +214,7 @@ def train_once(
   _reset_cuda_mem()
   data_rng, opt_init_rng, model_init_rng, rng = prng.split(rng, 4)
 
-  if FLAGS.framework == 'jax':
-    data_rng = jax.random.key_data(data_rng)
+  data_rng = jax.random.key_data(data_rng)
   # Workload setup.
   logging.info('Initializing dataset.')
   if hasattr(workload, '_eval_num_workers'):
@@ -348,8 +347,7 @@ def train_once(
     data_select_rng, update_rng, prep_eval_rng, eval_rng = \
       prng.split(step_rng, 4)
 
-    if FLAGS.framework == 'jax':
-      eval_rng = jax.random.key_data(eval_rng)
+    eval_rng = jax.random.key_data(eval_rng)
 
     with profiler.profile('Data selection'):
       batch = data_selection(workload,