Solved by ptrblck in post #2 
              
              
                Maybe not the most beautiful approach, but should get the work done: 
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc = nn.Linear(10, 10)
        self.non_fc = nn.Linear(1, 1)
        
    def forward(self, x):
        return x

def weight_in…