add __main__

phoBERT.py

@@ -1,6 +1,7 @@
 import torch
 from transformers import AutoModel, AutoTokenizer
 from underthesea import word_tokenize
+import __main__
 
 #phobert = AutoModel.from_pretrained("vinai/phobert-base")
 tokenizer = AutoTokenizer.from_pretrained("./")
@@ -31,6 +32,7 @@ class PhoBertModel(torch.nn.Module):
     return output
   
 def getModel():
+    __main__.PhoBertModel = PhoBertModel
     model = torch.load('phoBertModel.pth', map_location=torch.device('cpu'))
     model.eval()
     return model