Solved by ptrblck in post #2 
              
              
                The indexing of the input data and “selectively” calling backward works as expected in this code snippet: 
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.base = nn.Linear(1, 10)
        self.head1 = nn.Linear(10, 10)
        self.head2 = nn.Lin…