From d7f4605935614150863e1aa8508702599b28601a Mon Sep 17 00:00:00 2001 From: Sachin Prasad Date: Mon, 1 Apr 2024 15:36:30 -0700 Subject: [PATCH] Fix print_fn for custom function (#19419) --- keras/utils/summary_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/keras/utils/summary_utils.py b/keras/utils/summary_utils.py index 8b0d591d154..2fa9c7a0919 100644 --- 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):