Fix print_fn for custom function (#19419)

keras-team · Apr 1, 2024 · d7f4605 · d7f4605
1 parent 6c0d75f
commit d7f4605
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/keras/utils/summary_utils.py b/keras/utils/summary_utils.py
@@ -360,7 +360,10 @@ def print_layer(layer, nested_level=0):
 
     # Output captured summary for non-interactive logging.
     if print_fn:
-        print_fn(console.end_capture(), line_break=False)
+        if print_fn is io_utils.print_msg:
+            print_fn(console.end_capture(), line_break=False)
+        else:
+            print_fn(console.end_capture())
 
 
 def get_layer_index_bound_by_layer_name(layers, layer_range=None):