428516056abe41f135133e732a8d44af6ce9a234,rllib/policy/torch_policy.py,TorchPolicy,learn_on_batch,#TorchPolicy#Any#,221

Before Change

                    p.grad /= self.distributed_world_size
            info["allreduce_latency"] = time.time() - start

        self._optimizer.step()

        info.update(self.extra_grad_info(train_batch))
        return {
            LEARNER_STATS_KEY: info

After Change

            // Erase gradients in all vars of this optimizer.
            opt.zero_grad()
            // Recompute gradients of loss over all variables.
            loss_out[i].backward(retain_graph=(i < len(self._optimizers) - 1))

            grad_info.update(self.extra_grad_process(opt, loss_out[i]))

Instances