Reformatted code

Author: DimiChatzipavlis

keras/src/callbacks/__init__.py

@@ -14,4 +14,3 @@
 from keras.src.callbacks.swap_ema_weights import SwapEMAWeights
 from keras.src.callbacks.tensorboard import TensorBoard
 from keras.src.callbacks.terminate_on_nan import TerminateOnNaN
-

keras/src/callbacks/memory_usage_callback.py

@@ -1,14 +1,16 @@
 import os
 import warnings
+
+from keras.src import backend as K
 from keras.src.api_export import keras_export
 from keras.src.callbacks.callback import Callback
-from keras.src import backend as K
 
 # Attempt to import psutil for memory monitoring
 try:
     import psutil
 except ImportError:
-    psutil = None  
+    psutil = None
+
 
 @keras_export("keras.callbacks.MemoryUsageCallback")
 class MemoryUsageCallback(Callback):
@@ -63,14 +65,15 @@ def __init__(
 
         if tensorboard_log_dir:
             try:
-                import tensorflow as tf  
+                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}")
             except Exception as e:
                 warnings.warn(
-                    f"Could not initialize TensorBoard writer: {e}", RuntimeWarning
+                    f"Could not initialize TensorBoard writer: {e}",
+                    RuntimeWarning,
                 )
                 self._writer = None
 
@@ -85,7 +88,9 @@ 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):
@@ -122,7 +127,7 @@ def _get_gpu_memory(self):
         backend_name = K.backend()
         try:
             if backend_name == "tensorflow":
-                import tensorflow as tf  
+                import tensorflow as tf
 
                 gpus = tf.config.list_physical_devices("GPU")
                 if not gpus:
@@ -173,6 +178,8 @@ def _get_gpu_memory(self):
             return None
         except Exception as exc:
             if not hasattr(self, "_warn_exc"):
-                warnings.warn(f"Error retrieving GPU memory: {exc}", RuntimeWarning)
+                warnings.warn(
+                    f"Error retrieving GPU memory: {exc}", RuntimeWarning
+                )
                 self._warn_exc = True
             return None

keras/src/callbacks/memory_usage_callback_test.py

@@ -1,20 +1,21 @@
-import os
 import glob
+import os
 import re
 import sys
 import tempfile
-import pytest
-import numpy as np
-
 from contextlib import redirect_stdout
-from io import StringIO
 from importlib import reload
-from unittest.mock import patch, MagicMock
+from io import StringIO
+from unittest.mock import MagicMock
+from unittest.mock import patch
 
-from keras.src.models import Sequential
+import numpy as np
+import pytest
+
+from keras.src.callbacks.memory_usage_callback import MemoryUsageCallback
 from keras.src.layers import Dense
+from keras.src.models import Sequential
 from keras.src.testing import TestCase
-from keras.src.callbacks.memory_usage_callback import MemoryUsageCallback
 
 try:
     import psutil
@@ -69,13 +70,19 @@ def test_batch_logging_stdout(self):
         with redirect_stdout(buf):
             cb = MemoryUsageCallback(monitor_gpu=False, log_every_batch=True)
             self.model.fit(
-                self.x, self.y, epochs=1, batch_size=self.bs, callbacks=[cb], verbose=0
+                self.x,
+                self.y,
+                epochs=1,
+                batch_size=self.bs,
+                callbacks=[cb],
+                verbose=0,
             )
         lines = buf.getvalue().splitlines()
         batch_lines = [l for l in lines if l.startswith("Batch ")]
         assert len(batch_lines) == self.steps
         assert all(
-            re.match(r"Batch \d+ end - CPU Memory: [\d\.]+ MB", l) for l in batch_lines
+            re.match(r"Batch \d+ end - CPU Memory: [\d\.]+ MB", l)
+            for l in batch_lines
         )
 
     @pytest.mark.requires_trainable_backend
@@ -85,9 +92,16 @@ def test_tensorboard_writes_files(self):
         logdir = os.path.join(tmp.name, "tb")
         buf = StringIO()
         with redirect_stdout(buf):
-            cb = MemoryUsageCallback(monitor_gpu=False, tensorboard_log_dir=logdir)
+            cb = MemoryUsageCallback(
+                monitor_gpu=False, tensorboard_log_dir=logdir
+            )
             self.model.fit(
-                self.x, self.y, epochs=1, batch_size=self.bs, callbacks=[cb], verbose=0
+                self.x,
+                self.y,
+                epochs=1,
+                batch_size=self.bs,
+                callbacks=[cb],
+                verbose=0,
             )
         files = glob.glob(os.path.join(logdir, "events.out.tfevents.*"))
         assert files, "No TensorBoard event files generated"
@@ -118,7 +132,10 @@ def test_torch_backend_gpu_memory(monkeypatch):
     fake_torch = MagicMock()
     fake_torch.cuda.is_available.return_value = True
     fake_torch.cuda.device_count.return_value = 2
-    fake_torch.cuda.memory_allocated.side_effect = [100 * 1024**2, 150 * 1024**2]
+    fake_torch.cuda.memory_allocated.side_effect = [
+        100 * 1024**2,
+        150 * 1024**2,
+    ]
     monkeypatch.setitem(sys.modules, "torch", fake_torch)
 
     cb = MemoryUsageCallback(monitor_gpu=True)