Format the code

Author: DimiChatzipavlis

keras/src/callbacks/memory_usage_callback.py

@@ -19,11 +19,11 @@ class MemoryUsageCallback(Callback):
     This callback measures:
 
       - **CPU**: via psutil.Process().memory_info().rss
-      - **GPU/TPU**: via backend‐specific APIs
-      (TensorFlow, PyTorch, JAX, OpenVINO)
+      - **GPU/TPU**: via backend-specific APIs
+        (TensorFlow, PyTorch, JAX, OpenVINO)
 
     Logs are printed to stdout at the start/end of each epoch and,
-    if `log_every_batch=True`, after every batch.  If `tensorboard_log_dir`
+    if `log_every_batch=True`, after every batch. If `tensorboard_log_dir`
     is provided, scalars are also written via `tf.summary` (TensorBoard).
 
     Args:
@@ -36,7 +36,6 @@ class MemoryUsageCallback(Callback):
         ImportError: If `psutil` is not installed (required for CPU logging).
 
     Example:
-
     ```python
     from keras.callbacks import MemoryUsageCallback
     # ...
@@ -46,7 +45,7 @@ class MemoryUsageCallback(Callback):
         tensorboard_log_dir="./logs/memory"
     )
     model.fit(X, y, callbacks=[cb])
-    ```
+    ```  
     """
 
     def __init__(
@@ -67,7 +66,6 @@ def __init__(
         if tensorboard_log_dir:
             try:
                 import tensorflow as tf
-
                 logdir = os.path.expanduser(tensorboard_log_dir)
                 self._writer = tf.summary.create_file_writer(logdir)
                 print(f"MemoryUsageCallback: TensorBoard logs → {logdir}")
@@ -89,9 +87,7 @@ def on_epoch_end(self, epoch, logs=None):
 
     def on_batch_end(self, batch, logs=None):
         if self.log_every_batch:
-            self._log_step(
-                f"Batch {self._step_counter} end", self._step_counter
-            )
+            self._log_step(f"Batch {self._step_counter} end", self._step_counter)
         self._step_counter += 1
 
     def on_train_end(self, logs=None):
@@ -110,65 +106,66 @@ def _log_step(self, label, step):
         msg = f"{label} - CPU Memory: {cpu_mb:.2f} MB"
         if gpu_mb is not None:
             msg += f"; GPU Memory: {gpu_mb:.2f} MB"
-        print(msg)
+        # newline + flush ensures clean, immediate output
+        print("\n" + msg, flush=True)
 
         if self._writer:
             import tensorflow as tf  # noqa: E501
-
             with self._writer.as_default(step=int(step)):
                 tf.summary.scalar("Memory/CPU_MB", cpu_mb)
                 if gpu_mb is not None:
                     tf.summary.scalar("Memory/GPU_MB", gpu_mb)
             self._writer.flush()
 
     def _get_cpu_memory(self):
-        return self._proc.memory_info().rss / (1024**2)
+        return self._proc.memory_info().rss / (1024 ** 2)
 
     def _get_gpu_memory(self):
         backend_name = K.backend()
         try:
             if backend_name == "tensorflow":
                 import tensorflow as tf
-
                 gpus = tf.config.list_physical_devices("GPU")
                 if not gpus:
                     return None
                 total = sum(
                     tf.config.experimental.get_memory_info(g.name)["current"]
                     for g in gpus
                 )
-                return total / (1024**2)
+                return total / (1024 ** 2)
+
             elif backend_name == "torch":
                 import torch
-
                 if not torch.cuda.is_available():
                     return None
                 total = sum(
                     torch.cuda.memory_allocated(i)
                     for i in range(torch.cuda.device_count())
                 )
-                return total / (1024**2)
+                return total / (1024 ** 2)
+
             elif backend_name == "jax":
                 import jax
-
                 devs = [d for d in jax.devices() if d.platform.upper() == "GPU"]
                 if not devs:
                     return None
                 total = 0
                 for d in devs:
                     stats = getattr(d, "memory_stats", lambda: {})()
                     total += stats.get("bytes_in_use", 0)
-                return total / (1024**2)
-            else:
-                # OpenVINO and others fall back to unsupported
+                return total / (1024 ** 2)
 
+            else:
+                # OpenVINO and other unknown backends: warn once
                 if not hasattr(self, "_warn_backend"):
                     warnings.warn(
-                        f"MemoryUsageCallback: unsupported backend '{backend_name}'",
+                        "MemoryUsageCallback: unsupported backend "
+                        f"'{backend_name}'",
                         RuntimeWarning,
                     )
                     self._warn_backend = True
                 return None
+
         except ImportError as imp_err:
             if not hasattr(self, "_warn_import"):
                 warnings.warn(
@@ -177,6 +174,7 @@ def _get_gpu_memory(self):
                 )
                 self._warn_import = True
             return None
+
         except Exception as exc:
             if not hasattr(self, "_warn_exc"):
                 warnings.warn(