fixed attribute error

Author: ajkv-google

src/maxtext/trainers/post_train/distillation/distillation_utils.py

@@ -664,6 +664,7 @@ def __init__(
 
     # Re-initialize internal Orbax manager with MaxText's Grain handler
     # pylint: disable=access-member-before-definition
+    # pytype: disable=attribute-error
     if self._checkpoint_manager is not None:
       root_directory = self._checkpoint_manager.directory
 
@@ -684,6 +685,7 @@ def __init__(
           item_handlers=item_handlers,
           options=options,
       )
+    # pytype: enable=attribute-error
     # pylint: enable=access-member-before-definition
 
   def save(

src/maxtext/trainers/post_train/distillation/train_distill.py

@@ -378,13 +378,11 @@ def _log_metrics(self, loss, step=None, additional_metrics=None, **kwargs):
     tflops_per_sec = None
     if step_time_delta is not None and step_time_delta > 0:
       tflops_per_sec = self._tflops_combined / step_time_delta
-      tflops_metrics.update(
-          {
-              "perf/per_device_tflops_per_sec": tflops_per_sec,
-              "perf/per_device_tflops_per_sec_student": self._tflops_student / step_time_delta,
-              "perf/per_device_tflops_per_sec_teacher": self._tflops_teacher / step_time_delta,
-          }
-      )
+      tflops_metrics.update({
+          "perf/per_device_tflops_per_sec": tflops_per_sec,
+          "perf/per_device_tflops_per_sec_student": self._tflops_student / step_time_delta,
+          "perf/per_device_tflops_per_sec_teacher": self._tflops_teacher / step_time_delta,
+      })
     for name, value in tflops_metrics.items():
       self.metrics_logger.log(self.metrics_prefix, name, value, self._mode, step)