load_model fix

2023-05-22 23:19:48 +02:00
parent 14e78cf495
commit 08bf60a76a
1 changed files with 2 additions and 2 deletions
@@ -37,7 +37,7 @@ def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=No
    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'])
-    if scaler is not None: scaler.load_state_dict(checkpoint['scaler'])
+    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'])
    return int(iteration) if iteration != '' else 0