Spaces:

ki33elev
/

Transformer_arxiv_classification

Runtime error

ki33elev commited on Mar 26, 2022

Commit

e62c2f3

•

1 Parent(s): 45ac761

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -11,7 +11,7 @@ def load_model():
     model_name = 'distilbert-base-cased'
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=8)
-    model.load_state_dict(torch.load('model_weights.pt', map_location=torch.device('cpu')))
     model.eval()
     return tokenizer, model

     model_name = 'distilbert-base-cased'
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=8)
+    model.load_state_dict(torch.load('model_weights2.pt', map_location=torch.device('cpu')))
     model.eval()
     return tokenizer, model