Solved by ptrblck in post #2 
              
              
                You could use the inputs argument in backward to specify the gradient calculation for these parameters: 
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.first_part = nn.Linear(10, 10)
        self.second_part = nn.Linear(10, 10)

    def forward(self, input):
…