andresgtn commited on
Commit
8d3d679
β€’
1 Parent(s): e505593

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -9,14 +9,14 @@ extractor = AutoFeatureExtractor.from_pretrained("andresgtn/vit-base-bean-health
9
  model = AutoModelForImageClassification.from_pretrained("andresgtn/vit-base-bean-health-classifier")
10
 
11
  # add to cuda?
12
- model.eval()
13
- model.to(device)
14
 
15
  labels = dataset['train'].features['labels'].names
16
 
17
  def classify(im):
18
  features = extractor(im, return_tensors='pt')
19
- features.to(device) # move to gpu as model, if available
20
  with torch.no_grad():
21
  logits = model(**features).logits
22
  probability = torch.nn.functional.softmax(logits, dim=-1)
 
9
  model = AutoModelForImageClassification.from_pretrained("andresgtn/vit-base-bean-health-classifier")
10
 
11
  # add to cuda?
12
+ #model.eval()
13
+ #model.to(device)
14
 
15
  labels = dataset['train'].features['labels'].names
16
 
17
  def classify(im):
18
  features = extractor(im, return_tensors='pt')
19
+ #features.to(device) # move to gpu as model, if available
20
  with torch.no_grad():
21
  logits = model(**features).logits
22
  probability = torch.nn.functional.softmax(logits, dim=-1)