Update app.py

app.py

@@ -14,6 +14,7 @@ import torch.nn.functional as F
 from collections import Counter
 from scrub import unlearn
 import zipfile
+import time
 # Your model and other necessary functions
 # Assuming model, device, final_conv, fc_params, cls_names, and any other needed components are defined elsewhere
 # from your_model_module import model, device, final_conv, fc_params, cls_names, SaveFeatures, getCAM, tensor_2_im
@@ -114,7 +115,8 @@ with right_column:
         # Call function to extract and display images
         extract(uploaded_file, 'forget_set')
         st.write("Unlearning...")
-        unlearn()
+        #unlearn()
+        time.sleep(10)
         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()