upload weights only

predict.py

@@ -34,7 +34,9 @@ model_path = hf_hub_download(repo_id="Noha90/AML_16", filename="deit_best_model.
 print("Model path:", model_path)
 model = DeiT(num_classes=len(class_names))
 state_dict = torch.load(model_path, map_location="cpu")
-model.load_state_dict(state_dict)
+if "model_state_dict" in state_dict:
+    state_dict = state_dict["model_state_dict"]
+model.load_state_dict(state_dict, strict=False)
 model.eval()
 
 #deit transform