0ac2b33e8c63304a50db7d2b484368299706b58b,slm_lab/agent/net/recurrent.py,RecurrentNet,training_step,#RecurrentNet#Any#Any#Any#Any#Any#,164

Before Change

        if self.clip_grad:
            logger.debug(f"Clipping gradient: {self.clip_grad_val}")
            torch.nn.utils.clip_grad_norm_(self.parameters(), self.clip_grad_val)
        if global_net is None:
            self.optim.step()
        else:  // distributed training with global net
            net_util.push_global_grad(self, global_net)
            self.optim.step()
            net_util.pull_global_param(self, global_net)

After Change

    def training_step(self, x=None, y=None, loss=None, retain_graph=False, lr_t=None):
        """Takes a single training step: one forward and one backwards pass"""
        if hasattr(self, "model_tails") and x is not None:
            raise ValueError("Loss computation from x,y not supported for multitails")
        self.lr_scheduler.step(epoch=lr_t)
        self.train()
        self.optim.zero_grad()
        if loss is None:
            out = self(x)
            loss = self.loss_fn(out, y)
        assert not torch.isnan(loss).any(), loss
        if net_util.to_assert_trained():
            assert_trained = net_util.gen_assert_trained(self.rnn_model)
        loss.backward(retain_graph=retain_graph)
        if self.clip_grad_val is not None:
            nn.utils.clip_grad_norm_(self.parameters(), self.clip_grad_val)
        self.optim.step()
        if net_util.to_assert_trained():

Instances