From ded105b43eef1fc21a948d77b60e4cc0df18dcc8 Mon Sep 17 00:00:00 2001
From: Rino Lee <rinolee@google.com>
Date: Thu, 24 Nov 2022 05:11:00 -0800
Subject: [PATCH] Change PruningSummaries to only log pruning information

PiperOrigin-RevId: 490715299
---
 .../python/core/sparsity/keras/pruning_callbacks.py            | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensorflow_model_optimization/python/core/sparsity/keras/pruning_callbacks.py b/tensorflow_model_optimization/python/core/sparsity/keras/pruning_callbacks.py
index a51b3faa6..427d5ebf7 100644
--- a/tensorflow_model_optimization/python/core/sparsity/keras/pruning_callbacks.py
+++ b/tensorflow_model_optimization/python/core/sparsity/keras/pruning_callbacks.py
@@ -107,8 +107,7 @@ def _log_pruning_metrics(self, logs, prefix, step):
         self._file_writer.flush()
 
   def on_epoch_begin(self, epoch, logs=None):
-    if logs is not None:
-      super(PruningSummaries, self).on_epoch_begin(epoch, logs)
+    del logs
 
     pruning_logs = {}
     params = []