with class_weights

app.py

@@ -120,10 +120,11 @@ def load_model(test=False):
         with open(f'{path}/label_to_theme.json', 'r') as f:
             label_to_theme = json.load(f)
 
-    #class_weights = torch.load('model_info/class_weights.pth').to(device)
     # model = DebertPaperClassifier(device=device, num_labels=len(label_to_theme), class_weights=class_weights).to(device)
     # model.load_state_dict(torch.load("model_info/full_model_v4.pth", map_location=device))
-    model = DebertPaperClassifierV5(device=device, num_labels=47).to(device)
+
+    class_weights = torch.load(f'{path}/class_weights.pth').to(device)
+    model = DebertPaperClassifierV5(device=device, num_labels=47, class_weights=class_weights).to(device)
     model.load_state_dict(torch.load(f"{path}/full_model_v4.pth", map_location=device))
     if test:
         print(device)