Update app.py

app.py

@@ -121,6 +121,10 @@ with right_column:
         extract(uploaded_file, 'forget_set')
         st.write("Unlearning begins...")
         unlearn()
+        model_s = timm.create_model("rexnet_150", pretrained = True, num_classes = 17)  
+        model_s.load_state_dict(torch.load('celeb-model-unlearned.pth', map_location=torch.device('cpu')))
+        model_s.eval()
+
         uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
                 # Perform inference
         st.write("Performing inference...")
@@ -137,7 +141,7 @@ with right_column:
         preds = []
         with torch.no_grad():
             for i in range(50): 
-                output = model(image_tensor)
+                output = model_s(image_tensor)
                 probabilities = F.softmax(output, dim=1)
                 pred_class = torch.argmax(probabilities, dim=1)
                 pred_label = reversed_map[pred_class.item()]