Skip to content

GitLab

Unverified Commit 2215fa5c authored by Alexandros Koumparoulis's avatar Alexandros Koumparoulis Committed by GitHub
Browse files

don't put master_param to state if None (#1130) 


Signed-off-by: default avatarAlexandros Koumparoulis <akoumparouli@nvidia.com>
parent 309c6d49
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 3 deletions
transformer_engine/pytorch/optimizers/fused_adam.py
View file @ 2215fa5c
...@@ -219,10 +219,8 @@ class FusedAdam(torch.optim.Optimizer): ...@@ -219,10 +219,8 @@ class FusedAdam(torch.optim.Optimizer):
assert ( assert (
state["master_param"].shape == p.shape state["master_param"].shape == p.shape
), "Master weights shape must match model weights shape" ), "Master weights shape must match model weights shape"
else:
state["master_param"] = None
p_master = state["master_param"] p_master = state.get("master_param", None)
p_grad = p.grad p_grad = p.grad
if self.master_weights and p_master is not None and p_master.grad is not None: if self.master_weights and p_master is not None and p_master.grad is not None:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment