Textwizai / load_model.py
Erfan11's picture
Create load_model.py
0e7cada verified
raw
history blame
210 Bytes
import torch
from transformers import AutoModelForSequenceClassification
def load_model(model_path):
model = AutoModelForSequenceClassification.from_pretrained(model_path)
model.eval()
return model