Fix API integration

Author: DimiChatzipavlis

keras/api/_tf_keras/keras/callbacks/__init__.py

@@ -16,6 +16,9 @@
 from keras.src.callbacks.learning_rate_scheduler import (
     LearningRateScheduler as LearningRateScheduler,
 )
+from keras.src.callbacks.memory_usage_callback import (
+    MemoryUsageCallback as MemoryUsageCallback,
+)
 from keras.src.callbacks.model_checkpoint import (
     ModelCheckpoint as ModelCheckpoint,
 )

keras/api/callbacks/__init__.py

@@ -16,6 +16,9 @@
 from keras.src.callbacks.learning_rate_scheduler import (
     LearningRateScheduler as LearningRateScheduler,
 )
+from keras.src.callbacks.memory_usage_callback import (
+    MemoryUsageCallback as MemoryUsageCallback,
+)
 from keras.src.callbacks.model_checkpoint import (
     ModelCheckpoint as ModelCheckpoint,
 )

keras/src/callbacks/memory_usage_callback.py

@@ -19,7 +19,8 @@ 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`