update deit

predict.py

@@ -33,8 +33,8 @@ class DeiT(nn.Module):
 model_path = hf_hub_download(repo_id="Noha90/AML_16", filename="best_model.pth")
 print("Model path:", model_path)
 model = DeiT(num_classes=len(class_names))
-checkpoint = torch.load(model_path, map_location="cpu")
-model.load_state_dict(checkpoint["model_state_dict"])
+state_dict = torch.load(model_path, map_location="cpu")
+model.load_state_dict(state_dict)
 model.eval()
 
 #deit transform