diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py
index 097742c..0f718d5 100644
--- a/custom_utils/custom_utils.py
+++ b/custom_utils/custom_utils.py
@@ -291,22 +291,6 @@ def sparse_to_tensor(L):
     return tf.SparseTensor(idx, data, shape)
 
 
-def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False, iteration=''):
-    ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
-    if not disable_assert:
-        assert manager.latest_checkpoint is not None, checkpoint
-        assert manager.latest_checkpoint.split('/')[-1] != 'ckpt-0', checkpoint
-    if assert_consumed_on:
-        ckpt.restore(manager.latest_checkpoint).assert_consumed()
-    else:
-        if iteration != '':
-            name_base = manager.latest_checkpoint.split('-')[0]
-            ckpt_ind = [i for i in range(len(manager.checkpoints)) if manager.checkpoints[i].endswith(name_base + '-' + iteration)][0]
-            ckpt.restore(manager.checkpoints[ckpt_ind]).assert_existing_objects_matched()#.assert_consumed()
-        else:
-            ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
-
 
 def get_summary_str(model):
     # from torchview import draw_graph
diff --git a/utils/checkpoints.py b/utils/checkpoints.py
index 8ba4c5c..1699ada 100644
--- a/utils/checkpoints.py
+++ b/utils/checkpoints.py
@@ -5,7 +5,7 @@ import logging
 
 
 class CheckpointSaver:
-    def __init__(self, dirpath, decreasing=False, top_n=5):
+    def __init__(self, dirpath, top_n=5):
         """
         dirpath: Directory path where to store all model weights 
         decreasing: If decreasing is `True`, then lower metric is better
@@ -13,26 +13,28 @@ class CheckpointSaver:
         """
         if not os.path.exists(dirpath): os.makedirs(dirpath)
         self.dirpath = dirpath
-        self.top_n = top_n 
-        self.decreasing = decreasing
+        self.top_n = top_n
         self.top_model_paths = []
-        self.best_metric_val = np.Inf if decreasing else -np.Inf
         
-    def __call__(self, model, epoch, optimizer, metric_val=np.Inf):
-        model_path = os.path.join(self.dirpath, model.__class__.__name__ + f'_epoch{epoch}.pt')
-        save = metric_val<self.best_metric_val if self.decreasing else metric_val>self.best_metric_val
-        if save: 
-            logging.info(f"Current metric value better than {metric_val} better than best {self.best_metric_val}, saving model at {model_path}")
-            self.best_metric_val = metric_val
-            torch.save({ 'model': model.state_dict(), 'optimizer':optimizer }, model_path)
-            self.top_model_paths.append({'path': model_path, 'score': metric_val})
-            self.top_model_paths = sorted(self.top_model_paths, key=lambda o: o['score'], reverse=not self.decreasing)
+    def __call__(self, model, iteration, optimizer):
+        model_path = os.path.join(self.dirpath, f'iteration_{iteration}.pt')
+        torch.save({ 'model': model.state_dict(), 'optimizer': optimizer }, model_path)
+        self.top_model_paths.append({'path': model_path})
         if len(self.top_model_paths)>self.top_n: 
             self.cleanup()
     
     def cleanup(self):
-        to_remove = self.top_model_paths[self.top_n:]
+        to_remove = self.top_model_paths[:len(self.top_model_paths)-self.top_n]
         logging.info(f"Removing extra models.. {to_remove}")
         for o in to_remove:
             os.remove(o['path'])
-        self.top_model_paths = self.top_model_paths[:self.top_n]
+        self.top_model_paths = self.top_model_paths[len(self.top_model_paths)-self.top_n:]
+
+
+def load_model(model, checkpoint, optimizer=None, iteration=''):
+    # list of checkpoint files in the folder checkpoints
+    checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'{iteration}.pt')]
+    checkpoint = torch.load(checkpoints[-1])
+    model.load_state_dict(checkpoint['model'])
+    if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'])
+