quickfix

ui/model_loader.py

@@ -24,14 +24,27 @@ class ModelLoader:
 
         try:
             print(f"Attempting to fetch '{modelName}' from ClearML task: {taskID}")
-            
+
             task = Task.get_task(task_id=taskID)
 
-            modelObject = task.models["output"][0]
+            # Fetch the artifact 'model_one.pt'
+            artifact = task.artifacts.get(MODEL_ARTIFACT_NAME)
+
+            if artifact is None:
+                raise RuntimeError(
+                    f"Artifact '{MODEL_ARTIFACT_NAME}' not found in ClearML task {taskID}"
+                )
+
+            modelPath = artifact.get_local_copy()
+
+            if modelPath is None:
+                raise RuntimeError(
+                    f"Artifact '{MODEL_ARTIFACT_NAME}' could not be downloaded (returned None)"
+                )
 
-            modelPath = modelObject.get_local_copy()
             print(f"Weights downloaded to: {modelPath}")
 
+            # Load PyTorch model
             model = modelOne(noOfClasses=39)
             stateDict = torch.load(modelPath, map_location=self.device)
             model.load_state_dict(stateDict)
@@ -40,10 +53,11 @@ class ModelLoader:
             model.eval()
 
             return model
-        
+
         except Exception as e:
             print(f"Error loading from ClearML for {modelName}: {e}")
             raise RuntimeError(f"Failed to load model from ClearML: {e}")
+
     
     def loadModel(self, modelName):
         if modelName in self.modelCache: