Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 import torch
 from torchvision import transforms
+from efficientnet_pytorch import EfficientNet
 from huggingface_hub import HfFileSystem
 from PIL import Image
 
@@ -15,7 +16,8 @@ with open('best_model.pth', 'wb') as f:
     f.write(model_content)
 
 # Load your custom model onto the CPU
-model = torch.hub.load('ultralytics/yolov5', 'custom', path='best_model.pth', force_reload=True)
+model = EfficientNet.from_pretrained('efficientnet-b3')
+model.load_state_dict(torch.load('best_model.pth', map_location=torch.device('cpu')))
 model.eval()
 
 # Define a function that takes an image as input and uses the model for inference