Update unified_cell_classifier.py

unified_cell_classifier.py

@@ -250,12 +250,12 @@ class UnifiedCellClassifier(nn.Module):
         model.main_classifier = model._create_classifier_from_config(main_config)
         
         # Carica i pesi del main classifier
-        main_weights_path = get_file_path("main_classifier.bin")
-        main_state_dict = torch.load(main_weights_path, map_location="cpu")
+        main_weights_path = get_file_path("main_classifier/main_classifier.bin")
+        main_state_dict = torch.load(main_weights_path)
         model.main_classifier.load_state_dict(main_state_dict)
         
         # Carica le label del main classifier
-        main_labels_path = get_file_path("id2label_main.json")
+        main_labels_path = get_file_path("main_classifier/id2label_main.json")
         with open(main_labels_path) as f:
             model.main_labels = json.load(f)
         
@@ -271,7 +271,7 @@ class UnifiedCellClassifier(nn.Module):
                 
                 # Carica i pesi del sub-classificatore
                 sub_weights_path = get_file_path(f"sub_classifiers/{sub_name}.bin")
-                sub_state_dict = torch.load(sub_weights_path, map_location="cpu") 
+                sub_state_dict = torch.load(sub_weights_path) 
                 model.sub_classifiers[sub_name].load_state_dict(sub_state_dict)
                 
                 # Carica le label del sub-classificatore