How to implement local gradient in a tensor

[xiaolang01]

I want to achieve such an operation how the code should be written？

[akihironitta]

You should be able to achieve it by overriding on_after_backward with the logic. Something like this:

def on_after_backward(self):
    for p in self.model...parameters():
        p.grad[:original_vocab_size, :] = 0.0

https://pytorch-lightning.readthedocs.io/en/1.6.4/common/lightning_module.html#on-after-backward