Regarding BERT Arch during model deployment

I am begginer in NLP Transformers. I am facing this issue while deploying model using Django framework.Locally model is working fine but not when deployed. Here I am importing BERT model which was trained and saved using pytorch same procedure i follow to load the model but before I am defining the architecture which was defined during model training.

But facing the issue after deploying the model.

AttributeError: Can't get attribute 'BERT_Arch' on <module 'main' from '/home/ubuntu/kc_env/bin/gunicorn'>

I tried couple of things: Like defining the BERT Architecture before model loading:

####Utils.py

from django.apps import AppConfig
# import torch
import torch.nn as nn



class BERT_Arch(nn.Module):
    def __init__(self):
        super(BERT_Arch, self).__init__()
    
        # dropout layer
        self.dropout = nn.Dropout(0.2)
        
        # relu activation function
        self.relu =  nn.ReLU()
        # dense layer
        self.fc1 = nn.Linear(768,512)
        self.fc2 = nn.Linear(512,256)
        self.fc3 = nn.Linear(256,3)
        #softmax activation function
        self.softmax = nn.LogSoftmax(dim=1)
    #define the forward pass
    def forward(self, sent_id, mask):
        cls_hs = self.bert(sent_id, attention_mask=mask)[0][:,0]
    
        x = self.fc1(cls_hs)
        x = self.relu(x)
        x = self.dropout(x)
        
        x = self.fc2(x)
        x = self.relu(x)
        x = self.dropout(x)
        # output layer
        x = self.fc3(x)

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

###main.py

from .utils import BERT_Arch

model=BERT_Arch()

def func(): model=torch.load('Path to load model.pt')

Back to Top