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

load_model fix

Browse Source
This commit is contained in:
luni
2023-05-22 23:19:48 +02:00
parent 14e78cf495
commit 08bf60a76a
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/checkpoints.py
+2 -2
View File
@@ -37,7 +37,7 @@ def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=No
checkpoint = torch.load(checkpoints[-1]) checkpoint = torch.load(checkpoints[-1])
iteration = checkpoints[-1].split('_')[-1].split('.')[0] 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'].state_dict())
if scaler is not None: scaler.load_state_dict(checkpoint['scaler']) if scaler is not None and checkpoint['scaler'] is not None: scaler.load_state_dict(checkpoint['scaler'])
return int(iteration) if iteration != '' else 0 return int(iteration) if iteration != '' else 0