Saving the hidden layers of a trained network

I have used a pre-trained VGG-19 model to build an image classifier as part of a MOOC. I have implemented a classifier like this:

classifier = nn.Sequential(nn.Linear(25088, 512),
                           nn.ReLU(),
                           nn.Dropout(p=0.2),
                           nn.Linear(512,350),
                           nn.ReLU(),
                           nn.Dropout(p=0.2),
                           nn.Linear(350,250),
                           nn.ReLU(),
                           nn.Dropout(p=0.2),
                           nn.Linear(250,102),
                           nn.LogSoftmax(dim=1))

model.classifier = classifier

In the next step, I was asked to save the trained model so it could be reused later. My understanding of saving such a model was that I had to create a dictionary with the input size (given by the pre-trained model), the output size (given by my specific problem), the model’s state dict and the hidden layers (given by my own code). I tried to do that like this:

checkpoint = {'input_size':25088,
              'output_size':102,
              'hidden_layers':[each for each in model.classifier],
              'state_dict':model.state_dict(),
              'optimizer':optimizer.state_dict(),
              'epochs':epochs,
              'class_to_index':cat_to_name}

The I tried to load it like this:

def load_checkpoint(path):
    checkpoint = torch.load(path)
    model = nn.Sequential(checkpoint['input_size'],
                          checkpoint['output_size'],
                          checkpoint['hidden_layers'])
    model.load_state_dict(checkpoint['state_dict'])
    epochs = checkpoint['epochs']
    model.class_to_idx = checkpoint['class_to_index']
    return model, epochs

test_model, _ = load_checkpoint('trained_model.pth') 
print(test_model)

However, this just raises the error:

TypeError: int is not a Module subclass

What did I do wrong? As far as I can tell checkpoint[‘hidden_layers’] does not contain integers by themselves:

print(checkpoint['hidden_layers'])

Returns:

[Linear(in_features=25088, out_features=512, bias=True), ReLU(), Dropout(p=0.2), Linear(in_features=512, out_features=350, bias=True), ReLU(), Dropout(p=0.2), Linear(in_features=350, out_features=250, bias=True), ReLU(), Dropout(p=0.2), Linear(in_features=250, out_features=102, bias=True), LogSoftmax()]

Which is what I would expect.

Thank you for reading! Any thoughts or comments are highly appreciated!

def load_checkpoint(model, path):
    checkpoint = torch.load(path)
    model.classifier.load_state_dict(checkpoint['state_dict'])
    epochs = checkpoint['epochs']
    model.class_to_idx = checkpoint['class_to_index']
    return model, epochs

classifier = nn.Sequential(nn.Linear(25088, 512),
                           nn.ReLU(),
                           nn.Dropout(p=0.2),
                           nn.Linear(512,350),
                           nn.ReLU(),
                           nn.Dropout(p=0.2),
                           nn.Linear(350,250),
                           nn.ReLU(),
                           nn.Dropout(p=0.2),
                           nn.Linear(250,102),
                           nn.LogSoftmax(dim=1))

model.classifier = classifier

test_model, _ = load_checkpoint(model, 'trained_model.pth') 
print(test_model)

classifier = nn.Sequential(nn.Linear(25088, 512),
                           nn.ReLU(),
                           nn.Dropout(p=0.2),
                           nn.Linear(512,350),
                           nn.ReLU(),
                           nn.Dropout(p=0.2),
                           nn.Linear(350,250))

model.classifier = classifier
weights = torch.load("weight.pth")
state_dict = model.classifier.state_dict()

for name, param in model.classifier.named_parameters():
    if name in state_dict.keys():
        state_dict[name] = weights["state_dict"][name]

model.classifier.load_state_dict(state_dict)