mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix ET

Browse Source
This commit is contained in:
James Betker 2022-03-24 21:20:22 -06:00
parent a15970dd97
commit 07731d5491
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
codes/trainer/ExtensibleTrainer.py
View File

@ -318,7 +318,7 @@ class ExtensibleTrainer(BaseModel):
# This is needed to accurately log the grad norms.
for opt in step.optimizers:
from torch.cuda.amp.grad_scaler import OptState
if step.scaler._per_optimizer_states[id(opt)]["stage"] is not OptState.UNSCALED:
if step.scaler.is_enabled() and step.scaler._per_optimizer_states[id(opt)]["stage"] is not OptState.UNSCALED:
step.scaler.unscale_(opt)
if return_grad_norms and train_step: