set_to_none=True and accumulate_grad_batches

[denix56]

Is it possible to use automatic optimization with accumulate_grad_batches and performance trick zero_grad(set_to_none=True)?

[awaelchli]

Yes, you can override the zero_grad hook in LightningModule

    def optimizer_zero_grad(self, epoch: int, batch_idx: int, optimizer: Optimizer, optimizer_idx: int):
        optimizer.zero_grad(set_to_None=True)

[akihironitta]

https://pytorch-lightning.readthedocs.io/en/latest/benchmarking/performance.html#zero-grad-set-to-none-true