Help for adversarial learning with pytorch lighting #5690

fgm = FGM(model)
for batch_input, batch_label in data:
    # normal training
    loss = model(batch_input, batch_label)
    loss.backward() 
    # adversarial training
    fgm.attack() 
    loss_adv = model(batch_input, batch_label)
    loss_adv.backward() 
    fgm.restore()
    optimizer.step()
    model.zero_grad()

    def backward(self, loss: Tensor, optimizer: Optimizer, optimizer_idx: int, *args, **kwargs) -> None:
        if not self.hparams.do_adv:
            super().backward(loss, optimizer, optimizer_idx, *args, **kwargs)

    def training_step(self, batch, batch_idx):
        outputs = self(**batch)
        loss = outputs[0]
        if self.hparams.do_adv:
            loss.backward()
            self.fgm.attack()
            outputs = self(**batch)
            adv_loss = outputs[0]
            self.fgm.restore()
            adv_loss.backward()
        return loss