From 8255c1bbc94e761f3284590d35adaf932d99dec9 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sun, 21 May 2023 09:59:22 +0000 Subject: [PATCH] fix model loading --- utils/checkpoints.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/utils/checkpoints.py b/utils/checkpoints.py index 1699ada..c098f44 100644 --- a/utils/checkpoints.py +++ b/utils/checkpoints.py @@ -35,6 +35,8 @@ 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]) + iteration = checkpoints[-1].split('_')[-1].split('.')[0] model.load_state_dict(checkpoint['model']) if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer']) + return int(iteration) if iteration != '' else 0