From b6a1253a6e864b83a0394a7c2b6c375ef5210d1f Mon Sep 17 00:00:00 2001 From: Xibin Liu Date: Tue, 30 Dec 2025 22:45:55 +0000 Subject: [PATCH] Fix the LIBTPU_INIT_ARGS logging issue --- src/MaxText/metric_logger.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/MaxText/metric_logger.py b/src/MaxText/metric_logger.py index 77f1574b9b..9e3ace9bb0 100644 --- a/src/MaxText/metric_logger.py +++ b/src/MaxText/metric_logger.py @@ -273,7 +273,7 @@ def write_setup_info_to_tensorboard(self, params): self.metadata[MetadataKey.PER_DEVICE_TOKENS] = maxtext_utils.calculate_tokens_training_per_device(self.config) max_logging.log(f"number parameters: {num_model_parameters/1e9:.3f} billion") max_utils.add_text_to_summary_writer("num_model_parameters", str(num_model_parameters), self.writer) - max_utils.add_text_to_summary_writer("libtpu_init_args", os.environ["LIBTPU_INIT_ARGS"], self.writer) + max_utils.add_text_to_summary_writer("libtpu_init_args", os.getenv("LIBTPU_INIT_ARGS", ""), self.writer) maxtext_utils.add_config_to_summary_writer(self.config, self.writer) def get_performance_metric_queue(self, config):