diff --git a/baselines/jft/ood_utils.py b/baselines/jft/ood_utils.py
index 2ee823746..9f9bc4b20 100644
--- a/baselines/jft/ood_utils.py
+++ b/baselines/jft/ood_utils.py
@@ -435,8 +435,8 @@ def eval_ood_metrics(ood_ds,
 
         # Computes Maximum softmax probability (MSP)
         probs = jax.nn.softmax(logits[0], axis=-1)[masks_bool]
-        batch_scores['probs'] = probs
-        batch_scores['logits'] = logits[0][masks_bool]
+        batch_scores['probs'] = np.asarray(probs)
+        batch_scores['logits'] = np.asarray(logits)[0][masks_bool]
 
         # Compute Entropy
         batch_scores['entropy'] = np.array(
diff --git a/baselines/jft/test_utils.py b/baselines/jft/test_utils.py
index 11a7fd1d4..7d91a4720 100644
--- a/baselines/jft/test_utils.py
+++ b/baselines/jft/test_utils.py
@@ -114,7 +114,7 @@ def get_config(
 
   config.lr = ml_collections.ConfigDict()
   config.lr.base = 0.1
-  config.lr.warmup_steps = 1
+  config.lr.warmup_steps = 1 if total_steps > 1 else 0
   config.lr.decay_type = 'linear'
   config.lr.linear_end = 1e-3