add logging to loeading checkpoints

2023-07-02 14:59:46 +02:00
parent ef73cd1a20
commit b387e9d64d
1 changed files with 1 additions and 0 deletions
@@ -39,5 +39,6 @@ def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=No
    model.load_state_dict(checkpoint['model'])
    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'])
+    print(f'Loaded model from {checkpoints[-1]} at iteration {iteration}')
    return int(iteration) if iteration != '' else 0