Update app.py

app.py

@@ -45,10 +45,10 @@ def md_loading():
 
     print('ready')
 
-    return tokenizer, model, label_tbl, loc_tbl
+    return tokenizer, model, label_tbl, loc_tbl, device
 
 # 모델 로드
-tokenizer, model, label_tbl, loc_tbl = md_loading()
+tokenizer, model, label_tbl, loc_tbl, device = md_loading()
 
 
 # 데이터 셋 준비용
@@ -164,7 +164,7 @@ if st.button('확인'):
     # Predict 
     for batch in range(len(test_dataloader)):
         # Add batch to GPU
-#        batch = tuple(t.to(device) for t in batch)
+        batch = tuple(t.to(device) for t in batch)
 
         # Unpack the inputs from our dataloader
         test_input_ids, test_attention_mask = batch
@@ -178,7 +178,7 @@ if st.button('확인'):
         logits = outputs.logits
 
         # Move logits and labels to CPU
-        # logits = logits.detach().cpu().numpy()
+        logits = logits.detach().cpu().numpy()
 
 
     # # 단독 예측 시