Update app.py

app.py

@@ -45,21 +45,19 @@ process = psutil.Process(os.getpid())
 mem_before = process.memory_info().rss / (1024 * 1024)
 print(f"📦 모델 다운로드 전 메모리 사용량: {mem_before:.2f} MB")
 
+# 모델 로드
 try:
     model_path = hf_hub_download(repo_id=HF_MODEL_REPO_ID, filename=HF_MODEL_FILENAME)
     print(f"✅ 모델 파일 다운로드 성공: {model_path}")
 
-    mem_after_dl = process.memory_info().rss / (1024 * 1024)
-    print(f"📦 모델 다운로드 후 메모리 사용량: {mem_after_dl:.2f} MB")
-
-    # 모델 구성 및 state_dict 로드
-    model = CustomClassifier()
     state_dict = torch.load(model_path, map_location=device)
-    model.load_state_dict(state_dict)
+    model = BertForSequenceClassification.from_pretrained(
+        "skt/kobert-base-v1",
+        num_labels=len(category),
+        state_dict=state_dict,
+    )
+    model.to(device)
     model.eval()
-
-    mem_after_load = process.memory_info().rss / (1024 * 1024)
-    print(f"📦 모델 로드 후 메모리 사용량: {mem_after_load:.2f} MB")
     print("✅ 모델 로드 및 준비 완료.")
 except Exception as e:
     print(f"❌ Error: 모델 로드 중 오류 발생: {e}")