Update app.py

app.py

@@ -29,6 +29,9 @@ def predict(model, image_tensor):
     with torch.no_grad():
         outputs = model(image_tensor.unsqueeze(0))
         probs = torch.nn.functional.softmax(outputs[0], dim=0)
+        if torch.isnan(probs).any():
+            print("⚠️ Warning: NaN detected in prediction probabilities")
+            probs = torch.zeros_like(probs)
         pred = torch.argmax(probs).item()
     return probs, pred
  
@@ -41,9 +44,13 @@ def unlearn(model, image_tensor, label_idx, learning_rate, steps=10):
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.SGD(model.parameters(), lr=learning_rate)
  
-    for _ in range(steps):
+    for i in range(steps):
         output = model(image_tensor.unsqueeze(0))
         loss = -criterion(output, torch.tensor([label_idx]))
+        if torch.isnan(loss):
+            print(f"❌ NaN detected in loss at step {i}. Stopping unlearning.")
+            break
+        print(f"🧠 Step {i+1}/{steps} - Unlearning Loss: {loss.item():.4f}")
         optimizer.zero_grad()
         loss.backward()
         optimizer.step()
@@ -77,6 +84,7 @@ def run_unlearning(index_to_unlearn, learning_rate):
     # Get sample
     image_tensor, label_idx = trainset[index_to_unlearn]
     label_name = cifar10_classes[label_idx]
+    print(f"🗂️ Actual Label Index: {label_idx} | Label Name: {label_name}")
  
     # Prediction before
     probs_before, pred_before = predict(original_model, image_tensor)
@@ -94,7 +102,8 @@ def run_unlearning(index_to_unlearn, learning_rate):
     unlearn_acc, unlearn_loss = evaluate_model(unlearned_model, testloader)
  
     result = f"""
-📍 Index Unlearned: {index_to_unlearn} | Label: {label_name}
+📍 Index Unlearned: {index_to_unlearn}
+🗂️ Actual Label: {label_name} (Index: {label_idx})
  
 🔎 BEFORE Unlearning:
 - Prediction: {cifar10_classes[pred_before]}
@@ -117,7 +126,7 @@ demo = gr.Interface(
     fn=run_unlearning,
     inputs=[
         gr.Slider(0, len(trainset)-1, step=1, label="Select Index to Unlearn"),
-        gr.Slider(0.0001, 0.1, step=0.0001, value=0.01, label="Learning Rate (for Unlearning)")
+        gr.Slider(0.0001, 0.05, step=0.0001, value=0.005, label="Learning Rate (for Unlearning)")
     ],
     outputs="text",
     title="🔁 CIFAR-10 Machine Unlearning",