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

add logging to loeading checkpoints

Browse Source
This commit is contained in:
luni
2023-07-02 14:59:46 +02:00
parent ef73cd1a20
commit b387e9d64d
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/checkpoints.py
+1
View File
@@ -39,5 +39,6 @@ def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=No
model.load_state_dict(checkpoint['model']) model.load_state_dict(checkpoint['model'])
if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'].state_dict()) if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'].state_dict())
if scaler is not None and checkpoint['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'])
print(f'Loaded model from {checkpoints[-1]} at iteration {iteration}')
return int(iteration) if iteration != '' else 0 return int(iteration) if iteration != '' else 0