Update app.py

app.py

@@ -111,8 +111,8 @@ if (should_train_model=='1'): #train model
 
     
     # Model
-    model = BertForSequenceClassification.from_pretrained(repo_name, subfolder="bert_embeddings_finetune", output_attentions=True, num_labels=len(label_mapping), output_hidden_states=True).to('cuda')
-    # model = BertForSequenceClassification.from_pretrained('./mitra_ai_fleet_bert', output_attentions=True, num_labels=len(label_mapping), output_hidden_states=True).to('cuda')
+    model = BertForSequenceClassification.from_pretrained(repo_name, subfolder="bert_embeddings_finetune", output_attentions=True, num_labels=len(label_mapping), output_hidden_states=True).to('cpu')
+    # model = BertForSequenceClassification.from_pretrained('./mitra_ai_fleet_bert', output_attentions=True, num_labels=len(label_mapping), output_hidden_states=True).to('cpu')
     
 
     # Reset tokenizer size to include the new size after adding the tags to the tokenizer's tokens
@@ -353,7 +353,7 @@ else:
     model_save_path = "./saved_fleet_model"
     tokenizer_save_path = "./saved_fleet_tokenizer"
     # RobertaTokenizer.from_pretrained(model_save_path)
-    model = AutoModelForSequenceClassification.from_pretrained(model_save_path).to('cuda')
+    model = AutoModelForSequenceClassification.from_pretrained(model_save_path).to('cpu')
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_save_path)
 
 #Define the label mappings (this must match the mapping used during training)
@@ -370,7 +370,7 @@ def classify_user_input():
             break
              
         # Tokenize and predict
-        input_encoding = tokenizer(user_input, padding=True, truncation=True, return_tensors="pt").to('cuda')
+        input_encoding = tokenizer(user_input, padding=True, truncation=True, return_tensors="pt").to('cpu')
         
         with torch.no_grad():
             #attention_mask = input_encoding['attention_mask'].clone()