From 6e6b86cb4ad489c819758accfdf148128959ca2a Mon Sep 17 00:00:00 2001 From: Haohuan Wang Date: Wed, 21 Feb 2024 00:16:38 -0800 Subject: [PATCH] do not pull in tensorboard's keras callback --- clu/metric_writers/tf/summary_writer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/clu/metric_writers/tf/summary_writer.py b/clu/metric_writers/tf/summary_writer.py index 5711e70..436e519 100644 --- 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()