Load optimizer state from the checkpoint only when the training mode …

…is enabled (#632)
openvinotoolkit · Apr 5, 2021 · 83580b3 · 83580b3
1 parent f635ee6
commit 83580b3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/classification/staged_quantization_worker.py b/examples/classification/staged_quantization_worker.py
@@ -183,9 +183,9 @@ def autoq_eval_fn(model, eval_loader):
         best_acc1 = resuming_checkpoint['best_acc1']
         kd_loss_calculator.original_model.load_state_dict(resuming_checkpoint['original_model_state_dict'])
         compression_ctrl.scheduler.load_state(resuming_checkpoint['compression_scheduler'])
-        optimizer.load_state_dict(resuming_checkpoint['optimizer'])
-        optimizer_scheduler.load_state_dict(resuming_checkpoint['optimizer_scheduler'])
         if config.mode.lower() == 'train':
+            optimizer.load_state_dict(resuming_checkpoint['optimizer'])
+            optimizer_scheduler.load_state_dict(resuming_checkpoint['optimizer_scheduler'])
             logger.info("=> loaded checkpoint '{}' (epoch: {}, best_acc1: {:.3f})"
                         .format(resuming_checkpoint_path, resuming_checkpoint['epoch'], best_acc1))
         else: