Saving the hidden layers of a trained network
There are a few mistakes. The following is the way u need to change your code. I am not sure what your model class consists  but I am hoping it is class function with .classifier consisting the network.

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)


Similarly you need to load optimizer, scheduler and other attributes if you want to retrain the network.

if your intention is to load only part of the network instead of the entire network. Something as below

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)


In this way you can only load the part of the network required.