Solved by ptrblck in post #2 
              
              
                You could use a hook to zero out all other gradients. 
Here is a small example for a simple model: 
model = nn.Sequential(
    nn.Linear(2, 2),
    nn.Sigmoid(),
    nn.Linear(2, 2)
)
# Create Gradient mask
gradient_mask = torch.zeros(2, 2)
gradient_mask[0, 0] = 1.0
model[0].weight.register_hook(lam…