This fix adjusts import to avoid accidentally pulling in keras in clu. …

…#336 PiperOrigin-RevId: 617332775
google · Mar 21, 2024 · 580eaf0 · 580eaf0
1 parent c50acb7
commit 580eaf0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/clu/metric_writers/tf/summary_writer.py b/clu/metric_writers/tf/summary_writer.py
@@ -27,7 +27,7 @@
 from clu.metric_writers import interface
 import tensorflow as tf
 
-from tensorboard.plugins.hparams import api as hparams_api
+from tensorboard.plugins.hparams import summary_v2
 
 
 Array = interface.Array
@@ -95,7 +95,7 @@ def write_histograms(
 
   def write_hparams(self, hparams: Mapping[str, Any]):
     with self._summary_writer.as_default():
-      hparams_api.hparams(dict(utils.flatten_dict(hparams)))
+      summary_v2.hparams(dict(utils.flatten_dict(hparams)))
 
   def flush(self):
     self._summary_writer.flush()