Update app.py

app.py

@@ -1,17 +1,24 @@
-import requests
 import gradio as gr
-from fastai import *
+from fastai.vision.all import *
 
+# Train or load your fastai Learner object (replace this with your actual model training or loading code)
+path = untar_data(URLs.PETS)
+dls = ImageDataLoaders.from_folder(path/"images", valid_pct=0.2, seed=42)
+learn = cnn_learner(dls, resnet34)
+learn.fine_tune(1)
 
-# Load the trained model
-with open("model.pkl", "rb") as f:
-    learn = pickle.load(f)
+# Save the model using PyTorch's torch.save
+torch.save(learn.model, "model.pth")
 
 # Define the predict function
 def predict(image):
     # Load the input image
     img = PILImage.create(image)
 
+    # Load the model using PyTorch's torch.load
+    model = torch.load("model.pth")
+    model.eval()
+
     # Perform inference
     predicted_class, _, _ = learn.predict(img)