Anjali04-15 commited on
Commit
49876eb
·
verified ·
1 Parent(s): b827d72

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -3
app.py CHANGED
@@ -15,9 +15,11 @@ async def root():
15
  return {"message": "API is running"}
16
 
17
  # Load model and processor
18
- MODEL_NAME = "ivandrian11/fruit-classifier"
19
- processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
20
- model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
 
 
21
  model.eval() # set to evaluation mode
22
  DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
23
  model.to(DEVICE)
 
15
  return {"message": "API is running"}
16
 
17
  # Load model and processor
18
+ model_name = "ivandrian11/fruit-classifier"
19
+ AutoModelForImageClassification.from_pretrained(model_name).save_pretrained("fruit-model")
20
+ AutoImageProcessor.from_pretrained(model_name).save_pretrained("fruit-model")
21
+ processor = AutoImageProcessor.from_pretrained("fruit-model")
22
+ model = AutoModelForImageClassification.from_pretrained("fruit-model")
23
  model.eval() # set to evaluation mode
24
  DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
25
  model.to(DEVICE)