Fix formatting errors

Author: DimiChatzipavlis

keras/src/callbacks/__init__.py

@@ -6,11 +6,12 @@
 from keras.src.callbacks.history import History
 from keras.src.callbacks.lambda_callback import LambdaCallback
 from keras.src.callbacks.learning_rate_scheduler import LearningRateScheduler
+from keras.src.callbacks.memory_usage_callback import MemoryUsageCallback
 from keras.src.callbacks.model_checkpoint import ModelCheckpoint
 from keras.src.callbacks.progbar_logger import ProgbarLogger
 from keras.src.callbacks.reduce_lr_on_plateau import ReduceLROnPlateau
 from keras.src.callbacks.remote_monitor import RemoteMonitor
 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
-from keras.src.callbacks.memory_usage_callback import MemoryUsageCallback
+

keras/src/callbacks/memory_usage_callback.py

@@ -6,6 +6,7 @@
 from keras.src import backend as K
 
 # Attempt to import psutil for CPU memory
+
 try:
     import psutil
 except ImportError:
@@ -56,7 +57,7 @@ def __init__(
 
         if tensorboard_log_dir:
             try:
-                import tensorflow as tf  
+                import tensorflow as tf
 
                 logdir = os.path.expanduser(tensorboard_log_dir)
                 self.tb_writer = tf.summary.create_file_writer(logdir)
@@ -79,7 +80,7 @@ def _get_gpu_memory(self):
         backend = K.backend()
         try:
             if backend == "tensorflow":
-                import tensorflow as tf  
+                import tensorflow as tf
 
                 gpus = tf.config.list_physical_devices("GPU")
                 if not gpus:
@@ -89,9 +90,8 @@ def _get_gpu_memory(self):
                     info = tf.config.experimental.get_memory_info(gpu.name)
                     total += info.get("current", 0)
                 return total / (1024**2)
-
             if backend == "torch":
-                import torch  
+                import torch
 
                 if not torch.cuda.is_available():
                     return None
@@ -100,9 +100,8 @@ def _get_gpu_memory(self):
                     for i in range(torch.cuda.device_count())
                 )
                 return total / (1024**2)
-
             if backend == "jax":
-                import jax  
+                import jax
 
                 devs = [d for d in jax.devices() if d.platform == "gpu"]
                 if not devs:
@@ -112,15 +111,13 @@ def _get_gpu_memory(self):
                     stats = getattr(d, "memory_stats", lambda: {})()
                     total += stats.get("bytes_in_use", stats.get("allocated_bytes", 0))
                 return total / (1024**2)
-
             if not hasattr(self, "_warned_backend"):
                 warnings.warn(
                     f"Backend '{backend}' not supported for GPU memory.",
                     RuntimeWarning,
                 )
                 self._warned_backend = True
             return None
-
         except ImportError as e:
             warnings.warn(
                 f"Could not import backend lib ({e}); GPU disabled.",
@@ -139,7 +136,7 @@ def _log(self, label, step):
             msg += f"; GPU Memory: {gpu:.2f} MB"
         print(msg)
         if self.tb_writer:
-            import tensorflow as tf  
+            import tensorflow as tf
 
             with self.tb_writer.as_default(step=int(step)):
                 tf.summary.scalar("Memory/CPU_MB", cpu)

keras/src/callbacks/memory_usage_callback_test.py

@@ -1,7 +1,6 @@
 import os
 import glob
 import tempfile
-import warnings
 
 from contextlib import redirect_stdout
 from io import StringIO
@@ -18,6 +17,7 @@
 from keras.src import backend as K
 
 # Skip all tests if psutil is not installed
+
 try:
     import psutil
 except ImportError:
@@ -48,6 +48,7 @@ def test_epoch_and_batch_stdout(self):
         out = StringIO()
         with redirect_stdout(out):
             # Mock GPU memory for predictability
+
             with patch.object(
                 MemoryUsageCallback, "_get_gpu_memory", return_value=42.0
             ):
@@ -62,20 +63,24 @@ def test_epoch_and_batch_stdout(self):
                 )
         log = out.getvalue().splitlines()
         # Check epoch logs
+
         for i in range(self.epochs):
             assert any(f"Epoch {i} start" in line for line in log)
             assert any(f"Epoch {i} end" in line for line in log)
         # Check batch logs count
+
         batch_lines = [l for l in log if l.startswith("Batch")]
         assert len(batch_lines) == self.total_batches
         # Confirm GPU part present
+
         assert any("GPU Memory: 42.00 MB" in l for l in log)
 
     @pytest.mark.requires_trainable_backend
     def test_tensorboard_file_creation(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             tb_dir = os.path.join(tmpdir, "tb")
             # Mock CPU/GPU memory
+
             with patch.object(
                 MemoryUsageCallback, "_get_gpu_memory", return_value=10.0
             ), patch.object(MemoryUsageCallback, "_get_cpu_memory", return_value=5.0):
@@ -106,12 +111,15 @@ def test_import_error_without_psutil(self):
                 reload(mod)
                 _ = mod.MemoryUsageCallback()
         # restore
+
         if orig is not None:
             sys.modules["psutil"] = orig
             reload(mod)
 
 
 # Backend-specific tests
+
+
 @pytest.mark.requires_trainable_backend
 def test_torch_gpu_memory(monkeypatch):
     monkeypatch.setattr(K, "backend", lambda: "torch")