Load the checkpoint

checkpoint = torch.load(SAVE_PATH)

Reload the model

model = BertForSequenceClassification.from_pretrained(
    'bert-base-chinese', 
    num_labels=2,
    problem_type='single_label_classification'
    ).to(device)
model.load_state_dict(checkpoint['model_state_dict'])

Reload the optimizer

optimizer = AdamW(optimizer_grouped_parameters, lr=2e-5)
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

Other Info

test_acc = checkpoint['acc']  # corresponding test accuracy
epoch = checkpoint['epoch']  # which epoch the checkpoint belongs to