panda1835 commited on
Commit
c22649d
·
1 Parent(s): 44d687b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -0
app.py CHANGED
@@ -27,6 +27,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
27
 
28
  # Load the model
29
  classify_model = models.DinoVisionTransformerClassifier(num_classes)
 
30
  classify_model.load_state_dict(torch.load("best_dinov2_both_2023-11-21_07-44-35.pth", map_location=torch.device(device)))
31
  classify_model.eval()
32
 
 
27
 
28
  # Load the model
29
  classify_model = models.DinoVisionTransformerClassifier(num_classes)
30
+ classify_model = classify_model.to(device)
31
  classify_model.load_state_dict(torch.load("best_dinov2_both_2023-11-21_07-44-35.pth", map_location=torch.device(device)))
32
  classify_model.eval()
33