    def test_two_backward_callbacks(self):
        class SecondBackwardCallback(TrainerCallback):
            """
            Changes all gradients to 1 after backpropagation.
            """

            def on_backward(
                self,
                trainer: "GradientDescentTrainer",
                batch_outputs: Dict[str, torch.Tensor],
                backward_called: bool,
                **kwargs,
            ) -> bool:
                if backward_called:
                    raise OnBackwardException()
                batch_outputs["loss"].backward()
                for param in trainer.model.parameters():
                    param.grad = torch.ones_like(param.grad, device=param.grad.device)
                return True

        with pytest.raises(OnBackwardException):
            trainer = GradientDescentTrainer(
                self.model,
                self.optimizer,
                self.data_loader,
                num_epochs=2,
                validation_data_loader=self.validation_data_loader,
                callbacks=[
                    ZeroGradientsBackwardCallback(serialization_dir=self.TEST_DIR),
                    SecondBackwardCallback(serialization_dir=self.TEST_DIR),
                ],
            )
            trainer.train()
Changes all gradients to 1 after backpropagation.
tests/training/trainer_test.py