From 125d7234f1c49924fd2e3a54ac13c4464515279e Mon Sep 17 00:00:00 2001 From: Sachin Prasad Date: Mon, 1 Apr 2024 14:33:24 -0700 Subject: [PATCH] Fix print_fn for custom function --- 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):