fix model loading

2023-05-21 09:59:22 +00:00
parent 391e7f120b
commit 8255c1bbc9
1 changed files with 2 additions and 0 deletions
@@ -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