updating the code to remove older way of tree map

Author: init-22

algoperf/workloads/cifar/cifar_jax/workload.py

@@ -220,4 +220,4 @@ def _normalize_eval_metrics(
     self, num_examples: int, total_metrics: Dict[str, Any]
   ) -> Dict[str, float]:
     """Normalize eval metrics."""
-    return jax.tree_map(lambda x: x / num_examples, total_metrics)
+    return jax.tree.map(lambda x: x / num_examples, total_metrics)