Gather function in Lightning module with gradients
🚀 Feature
Add a gather function with gradients in LightningModule which can be implemented with the help of:
https://github.com/PyTorchLightning/pytorch-lightning-bolts/blob/master/pl_bolts/models/self_supervised/simclr/simclr_module.py#L25
Motivation
The default all_gather in pytorch breaks the gradient graph as it does not propagate any gradient backwards.
Pitch
Allow users to use self.gather in LightningModule.
x = self.gather(x)

As suggested by @williamFalcon, this can be implemented like:
def gather(self):
    self.trainer.gather():


Trainer.gather(self, x):
    x = self.gather.accelerator(x)

All accelerators except CPU should implement this.
The gradients can be synced with:
x = self.gather(x, sync_grads=True)

@edenlightning @tchaton @teddykoker