update app.py

app.py

@@ -1,4 +1,4 @@
-### 1. Imports and class names setup ### 
+### 1. Imports and class names setup ###
 import gradio as gr
 import os
 import torch
@@ -8,18 +8,18 @@ from timeit import default_timer as timer
 from typing import Tuple, Dict
 
 # Setup class names
-with open("class_names.txt", "r") as f: # reading them in from class_names.txt
-    class_names = [food_name.strip() for food_name in  f.readlines()]
-    
-### 2. Model and transforms preparation ###    
+with open("class_names.txt", "r") as f:  # reading them in from class_names.txt
+    class_names = [food_name.strip() for food_name in f.readlines()]
+
+### 2. Model and transforms preparation ###
 
 # Create model
-model, transforms = create_effnetb2_model(
-    num_classes=101, # could also use len(class_names)
+effnetb2, effnetb2_transforms = create_effnetb2_model(
+    num_classes=101,  # could also use len(class_names)
 )
 
 # Load saved weights
-model.load_state_dict(
+effnetb2.load_state_dict(
     torch.load(
         f="09_pretrained_effnetb2_feature_extractor_food101_20_percent.pth",
         map_location=torch.device("cpu"),  # load to CPU
@@ -34,25 +34,28 @@ def predict(img) -> Tuple[Dict, float]:
     """
     # Start the timer
     start_time = timer()
-    
+
     # Transform the target image and add a batch dimension
     img = effnetb2_transforms(img).unsqueeze(0)
-    
+
     # Put model into evaluation mode and turn on inference mode
     effnetb2.eval()
     with torch.inference_mode():
         # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
         pred_probs = torch.softmax(effnetb2(img), dim=1)
-    
+
     # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
-    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
-    
+    pred_labels_and_probs = {
+        class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))
+    }
+
     # Calculate the prediction time
     pred_time = round(timer() - start_time, 5)
-    
-    # Return the prediction dictionary and prediction time 
+
+    # Return the prediction dictionary and prediction time
     return pred_labels_and_probs, pred_time
 
+
 ### 4. Gradio app ###
 
 # Create title, description and article strings
@@ -63,7 +66,7 @@ article = "Created at [09. PyTorch Model Deployment](https://www.learnpytorch.io
 # Create examples list from "examples/" directory
 example_list = [["examples/" + example] for example in os.listdir("examples")]
 
-# Create Gradio interface 
+# Create Gradio interface
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),