try: state_dict

app.py

@@ -32,9 +32,17 @@ for i, url in enumerate(file_urls):
             f"image_{i}.jpg"
         )
 
-model = Model('models/custom_yolov5s.yaml')  
-model.load_state_dict(torch.load('ppe-yolov5.pt')) 
-model.eval()
+# Load the checkpoint
+checkpoint = torch.load('ppe-yolov5.pt')
+
+# Ensure 'model' is a key in the checkpoint before attempting to load the state_dict
+if 'model' in checkpoint:
+    model = Model('models/custom_yolov5s.yaml')
+    model.load_state_dict(checkpoint['model'])  # Assuming the key for state_dict is 'model'
+    model.eval()
+else:
+    raise KeyError("The checkpoint file does not contain the key 'model'. Please check the file.")
+
 path = [['image_0.jpg'], ['image_1.jpg']]
 video_path = [['video.mp4']]