fix for fp16 loss scaling (#52)

fixes #44
openai · Jul 15, 2022 · 22e0df8 · 22e0df8
1 parent 4afb0ac
commit 22e0df8
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/guided_diffusion/fp16_util.py b/guided_diffusion/fp16_util.py
@@ -199,7 +199,8 @@ def _optimize_fp16(self, opt: th.optim.Optimizer):
         logger.logkv_mean("grad_norm", grad_norm)
         logger.logkv_mean("param_norm", param_norm)
 
-        self.master_params[0].grad.mul_(1.0 / (2 ** self.lg_loss_scale))
+        for p in self.master_params:
+            p.grad.mul_(1.0 / (2 ** self.lg_loss_scale))
         opt.step()
         zero_master_grads(self.master_params)
         master_params_to_model_params(self.param_groups_and_shapes, self.master_params)