From 67176383c4195d3c13a91e78b6f8701d9d579dec Mon Sep 17 00:00:00 2001
From: Sylwester Klocek <sylwester.klocek@uj.edu.pl>
Date: Mon, 5 Oct 2020 21:23:18 +0200
Subject: [PATCH] Update engine.py (#458)

Do not load optimizer states for fp32 if flag set to False.
---
 deepspeed/runtime/engine.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 0b24595b..cbc4aa7e 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -1227,7 +1227,7 @@ class DeepSpeedEngine(Module):
                 self.optimizer.load_state_dict(
                     checkpoint['optimizer'],
                     load_optimizer_states=load_optimizer_states)
-            else:
+            elif load_optimizer_states:
                 self.optimizer.load_state_dict(checkpoint['optimizer'])
 
         if load_lr_scheduler_states and self.lr_scheduler is not None:
-- 
GitLab