luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix model loading

Browse Source
This commit is contained in:
luni
2023-05-21 09:59:22 +00:00
parent 391e7f120b
commit 8255c1bbc9
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/checkpoints.py
+2
View File
@@ -35,6 +35,8 @@ def load_model(model, checkpoint, optimizer=None, iteration=''):
# list of checkpoint files in the folder checkpoints # 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')] checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'{iteration}.pt')]
checkpoint = torch.load(checkpoints[-1]) checkpoint = torch.load(checkpoints[-1])
iteration = checkpoints[-1].split('_')[-1].split('.')[0]
model.load_state_dict(checkpoint['model']) model.load_state_dict(checkpoint['model'])
if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer']) if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'])
return int(iteration) if iteration != '' else 0