Regarding attribute error BERT Arch during model deployment using Django - django

Question

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')