snapshot...

modules/GSOC/E2_ESRGAN/lib/train.py

@@ -66,7 +66,6 @@ class Trainer(object):
     decay_step = decay_params["step"]
     decay_factor = decay_params["factor"]
     total_steps = phase_args["num_steps"]
-    print('total number of steps: ', total_steps)
     metric = tf.keras.metrics.Mean()
     psnr_metric = tf.keras.metrics.Mean()
 
@@ -79,6 +78,7 @@ class Trainer(object):
     checkpoint = tf.train.Checkpoint(
         G=generator,
         G_optimizer=G_optimizer)
+    ckpt_manager = tf.train.CheckpointManager(checkpoint, self.model_dir, max_to_keep=7)
 
     status = utils.load_checkpoint(checkpoint, "phase_1", self.model_dir)
     logging.debug("phase_1 status object: {}".format(status))