Merge branch 'main' of github.com:AI-Hypercomputer/maxtext into sujin…

…esh/llama2_v6e_pw_long_running_test
AI-Hypercomputer · Jan 9, 2025 · 8b2fc50 · 8b2fc50
2 parents e5114d4 + d01c2aa
commit 8b2fc50
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/MaxText/train.py b/MaxText/train.py
@@ -963,12 +963,13 @@ def train_loop(config, state=None):
     # pytype: disable=attribute-error
     compiled = p_train_step.lower(state, example_batch, nextrng).compile()
     compiled_stats = compiled.memory_analysis()
-    max_logging.log(
-        f"Output size: {compiled_stats.output_size_in_bytes}, "
-        f"temp size: {compiled_stats.temp_size_in_bytes}, "
-        f"argument size: {compiled_stats.argument_size_in_bytes}, "
-        f"host temp size: {compiled_stats.host_temp_size_in_bytes}, in bytes."
-    )
+    if compiled_stats is not None:
+      max_logging.log(
+          f"Output size: {compiled_stats.output_size_in_bytes}, "
+          f"temp size: {compiled_stats.temp_size_in_bytes}, "
+          f"argument size: {compiled_stats.argument_size_in_bytes}, "
+          f"host temp size: {compiled_stats.host_temp_size_in_bytes}, in bytes."
+      )
   return state