Internal logging change

PiperOrigin-RevId: 632260125 Change-Id: Idf7cf06cab660aa9f8e98926a48780bb7d3c5832
google-deepmind · May 9, 2024 · 290d13a · 290d13a
1 parent b89055b
commit 290d13a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/neural_testbed/logging/base.py b/neural_testbed/logging/base.py
@@ -88,10 +88,10 @@ def problem(self) -> testbed_base.TestbedProblem:
 
 
 def clean_results(results: Dict[str, Any]) -> Dict[str, Any]:
-  """Cleans the results for logging (can't log jax arrays)."""
+  """Cleans the results for logging."""
   def clean_result(value: Any) -> Any:
     value = loggers.to_numpy(value)
-    if isinstance(value, chex.ArrayNumpy) and value.size == 1:
+    if isinstance(value, chex.Array) and value.size == 1:
       value = float(value)
     return value