Noha90 commited on
Commit
248c755
·
1 Parent(s): 9937dde
Files changed (1) hide show
  1. predict.py +1 -1
predict.py CHANGED
@@ -31,7 +31,7 @@ class DeiT(nn.Module):
31
 
32
  # Load model
33
  model_path = hf_hub_download(repo_id="Noha90/AML_16", filename="deit_best_model.pth", repo_type="space")
34
- print("Model path_check:", model_path)
35
  model = DeiT(num_classes=len(class_names))
36
  state_dict = torch.load(model_path, map_location="cpu")
37
  model.load_state_dict(state_dict)
 
31
 
32
  # Load model
33
  model_path = hf_hub_download(repo_id="Noha90/AML_16", filename="deit_best_model.pth", repo_type="space")
34
+ print("Model path:", model_path)
35
  model = DeiT(num_classes=len(class_names))
36
  state_dict = torch.load(model_path, map_location="cpu")
37
  model.load_state_dict(state_dict)