Update main.py

ssu-Recipable · May 28, 2024 · b74581a · b74581a
1 parent 54efc48
commit b74581a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/source/main.py b/source/main.py
@@ -10,7 +10,7 @@
 import uvicorn
 
 tokenizer = BertTokenizer.from_pretrained('kykim/bert-kor-base')
-model = BertForSequenceClassification.from_pretrained('kykim/bert-kor-base', num_labels=317)
+model = BertForSequenceClassification.from_pretrained('kykim/bert-kor-base', num_labels=318)
 model.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
 
 def pad_sequences(input_ids, maxlen) :