Update app.py

app.py

@@ -95,7 +95,7 @@ option = st.selectbox("Select a text analysis model:", ("BERT", "Fine-tuned BERT
 if option == "BERT":
     model = PretrainedBertClass()
 else:
-    model = torch.load("pytorch_bert_toxic.bin")
+    model = torch.load("pytorch_bert_toxic.bin", map_location=torch.device("cpu"))
 
 # Freeze model and input tokens
 def inference():
@@ -108,7 +108,7 @@ def inference():
             mask = data["mask"].to(device, dtype=torch.long)
             token_type_ids = data["token_type_ids"].to(device, dtype=torch.long)
             targets = data["targets"].to(device, dtype=torch.float)
-            outputs = model(ids, mask, token_type_ids)
+            outputs = model(ids, mask, token_type_ids, return_dict=False)
             final_targets.extend(targets.cpu().detach().numpy().tolist())
             final_outputs.extend(torch.sigmoid(outputs).cpu().detach().numpy().tolist())
     return final_outputs, final_targets