60dfcf812eca79017dced46e1189245c050a3fd6,dl/callbacks.py,OptimizerCallback,on_batch_end,#OptimizerCallback#Any#,393

Before Change

                assert len(state._optimizer) == 1, \
                    "fp16 mode works only with one optimizer for now"

                for key, value in state.loss.items():
                    scaled_loss = self.fp16_grad_scale * value.float()
                    scaled_loss.backward()

                master_params = list(
                    state._optimizer["main"].param_groups[0]["params"])
                model_params = list(filter(
                    lambda p: p.requires_grad,

After Change

                    optimizer=optimizer,
                    optimizer_wd=self.optimizer_wd,
                    grad_clip_fn=self.grad_clip_fn)
                model.zero_grad()
                self.accumulation_counter = 0
        else:
            model = state.model
            model.zero_grad()

Instances