How to load a fine tuned pytorch huggingface bert model from a checkpoint file?

I had fine tuned a bert model in pytorch and saved its checkpoints via torch.save(model.state_dict(), 'model.pt')

Now When I want to reload the model, I have to explain whole network again and reload the weights and then push to the device.

Can anyone tell me how can I save the bert model directly and load directly to use in production/deployment?

Following is the training code and you can try running there in colab itself! After training completion, you will notice in file system we have a checkpoint file. But I want to save the model itself.

LINK TO COLAB NOTEBOOK FOR SAMPLE TRAINING

Following is the current inferencing code I written.

import torch
import torch.nn as nn
from transformers import AutoModel, BertTokenizerFast
import numpy as np
import json

tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
device = torch.device("cpu")


class BERT_Arch(nn.Module):

    def __init__(self, bert):
        super(BERT_Arch, self).__init__()

        self.bert = bert

        # dropout layer
        self.dropout = nn.Dropout(0.1)

        # relu activation function
        self.relu = nn.ReLU()

        # dense layer 1
        self.fc1 = nn.Linear(768, 512)

        # dense layer 2 (Output layer)
        self.fc2 = nn.Linear(512, 2)

        # softmax activation function
        self.softmax = nn.LogSoftmax(dim=1)

    # define the forward pass
    def forward(self, sent_id, mask):
        # pass the inputs to the model
        _, cls_hs = self.bert(sent_id, attention_mask=mask, return_dict=False)

        x = self.fc1(cls_hs)

        x = self.relu(x)

        x = self.dropout(x)

        # output layer
        x = self.fc2(x)

        # apply softmax activation
        x = self.softmax(x)

        return x


bert = AutoModel.from_pretrained('bert-base-uncased')
model = BERT_Arch(bert)
path = './models/saved_weights_new_data.pt'
model.load_state_dict(torch.load(path, map_location=device))
model.to(device)

def inference(comment):
    tokens_test = tokenizer.batch_encode_plus(
        list([comment]),
        max_length=75,
        pad_to_max_length=True,
        truncation=True,
        return_token_type_ids=False
    )
    test_seq = torch.tensor(tokens_test['input_ids'])
    test_mask = torch.tensor(tokens_test['attention_mask'])
    predictions = model(test_seq.to(device), test_mask.to(device))
    predictions = predictions.detach().cpu().numpy()
    predictions = np.argmax(predictions, axis=1)
    return predictions

I simply want to save a model from this notebook in a way such that I can use it for inferencing anywhere.