Update app.py

app.py

@@ -7,7 +7,7 @@ import re
 
 def paragraph_leveling(text):
     model_name = "./trained_model/fine_tunning_encoder_v2"
-    model = AutoModel.from_pretrained(model_name)
+    model = AutoModel.from_pretrained(model_name).to('cpu')
     tokenizer = AutoTokenizer.from_pretrained('zzxslp/RadBERT-RoBERTa-4m')
 
     class MLP(nn.Module):
@@ -23,6 +23,7 @@ def paragraph_leveling(text):
 
     classifier = MLP(target_size=3, input_size=768)
     classifier.load_state_dict(torch.load('./trained_model/fine_tunning_classifier'))
+    classifier.to('cpu')
     classifier.eval()
 
     output_list = []
@@ -44,7 +45,7 @@ def paragraph_leveling(text):
             padding='max_length',
             truncation=True,
             max_length=120)
-        output = model(**train_encoding)
+        output = model(**train_encoding.to('cpu'))
         output = classifier(output[1])
         output = output[0]