Solved by spanev in post #2 
              
              
                Hi @isalirezag, 
Do you want to apply it during back-propagation or after? 
In the first case, you should be able to achieve it with a backward hook: 
def hook_fn(module, grad_input, grad_output):
    if isinstance(grad_input, (list, tuple)):
        return tuple((grad + C if not grad is None else N…