Update app.py

app.py

@@ -184,8 +184,9 @@ if st.button('확인'):
 
         # Move logits and labels to CPU
 #        logits = logits.detach().cpu().numpy()
-
-    st.write(logits.size())
+    pred_m = torch.nn.Softmax(dim=1)
+    pred_ = pred_m(logits)
+    # st.write(logits.size())
     # # 단독 예측 시
     # arg_idx = torch.argmax(logits, dim=1)
     # print('arg_idx:', arg_idx)
@@ -195,15 +196,15 @@ if st.button('확인'):
 
     # 상위 k번째까지 예측 시
     k = 10
-    topk_idx = torch.topk(logits.flatten(), k).indices
-    topk_values = torch.topk(logits.flatten(), k).values
+    topk_idx = torch.topk(pred_.flatten(), k).indices
+    topk_values = torch.topk(pred_.flatten(), k).values
 
 
     num_ans_topk = label_tbl[topk_idx]
     str_ans_topk = [loc_tbl['항목명'][loc_tbl['코드'] == k] for k in num_ans_topk]
     percent_ans_topk = topk_values.numpy()
     
-    st.write(sum(torch.topk(logits.flatten(), 493).values.numpy()))
+    st.write(sum(torch.topk(pred_.flatten(), 493).values.numpy()))
     # print(num_ans, str_ans)
     # print(num_ans_topk)