quickfix variable name mixup

ui/app.py

@@ -19,7 +19,7 @@ from model_loader import ModelLoader
 class PlantDiseaseApp:
     def __init__(self):
         self.model_loader = ModelLoader()
-        self.current_modelName = "model_one.pt"
+        self.current_modelName = "CNN from Scratch"
         self.model = self.model_loader.loadModel(self.current_modelName)
         self.flagged_predictions = []
 

ui/model_loader.py

@@ -20,7 +20,6 @@ class ModelLoader:
             raise ValueError(f"ClearML configuration not found for model: {modelName}")
 
         taskID = modelConfig['clearml_task_id']
-        modelType  = modelConfig['model_type']
 
         try:
             print(f"attemtping to fetch '{modelName}' from clearML task: {taskID}")
@@ -29,10 +28,10 @@ class ModelLoader:
             modelPath = modelObject.get_local_copy()
             print(f"Weights downloaded to local path: {modelPath}")
 
-            model = modelOne()
+            model = modelOne(noOfClasses=39)
 
             stateDict = torch.load(modelPath, map_location=self.device)
-            modelObject.load_state_dict(stateDict)
+            model.load_state_dict(stateDict)
 
             model.to(self.device)
             model.eval()