bea3a486261d3f6238316224605b684db1e5bca7,Tars/models/custom_loss_model.py,CustomLossModel,train,#CustomLossModel#Any#,36

Before Change

    def train(self, train_x):
        self.distributions.train()

        self.optimizer.zero_grad()
        loss = self.loss_cls.estimate(train_x).mean()

        // backprop
        loss.backward()

        // update params
        self.optimizer.step()

        return loss

    def test(self, test_x):

After Change

    def train(self, train_x, **kwargs):
        self.distributions.train()

        return super().train(train_x, **kwargs)

    def test(self, test_x, **kwargs):
        self.distributions.eval()

Instances