Solved by ptrblck in post #2 
              
              
                I would save the whole model’s state_dict and just reimplement an “inference” model, which yields only one output. Here is a small example: 

class MyModelA(nn.Module):
    def __init__(self):
        super(MyModelA, self).__init__()
        self.fc1 = nn.Linear(10, 1)
        self.fc2 = nn.Linear(1…