diff --git a/utils/checkpoints.py b/utils/checkpoints.py
index ebbd87a..30db495 100644
--- a/utils/checkpoints.py
+++ b/utils/checkpoints.py
@@ -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