Update app.py

push the targets outside loop

app.py

@@ -74,29 +74,34 @@ def simple_generative_inference(image, mode, model, n_iterations=10, step_size=0
     image_tensor = transform(image).unsqueeze(0).requires_grad_(True)
     image_tensor.retain_grad()  # Ensure gradients are retained for non-leaf tensor
 
+    if mode == "increase confidence":
+        output = model(image_tensor)
+        probs = torch.nn.functional.softmax(output, dim=1)
+        _, least_likely_indices = torch.topk(probs, k=2, largest=False)
+    elif mode == "ReverseDiffuse":
+        noisy_image = image_tensor + torch.randn_like(image_tensor) * noise_ratio
+    else:
+        raise ValueError("Invalid mode selected. Choose 'increase confidence' or 'ReverseDiffuse'.")
+        
     for _ in range(n_iterations):
+
+        image_tensor
         # Zero gradients
-        if image_tensor.grad is not None:
-            image_tensor.grad.zero_()
+        model.zero_grad()
 
         # Forward pass
         output = model(image_tensor)
 
         # Define inference loss based on mode
         if mode == "increase confidence":
-            probs = torch.nn.functional.softmax(output, dim=1)
-            _, least_likely_indices = torch.topk(probs, k=2, largest=False)
             losses = []
             for idx in least_likely_indices[0]:
                 target = torch.full((1,), idx, dtype=torch.long, device=output.device)
                 loss = torch.nn.CrossEntropyLoss()(output, target)
                 losses.append(loss)
             loss = torch.stack(losses).mean()
-        elif mode == "ReverseDiffuse":
-            noisy_image = image_tensor + torch.randn_like(image_tensor) * noise_ratio
+        elif mode == "ReverseDiffuse":            
             loss = torch.nn.functional.mse_loss(image_tensor, noisy_image)
-        else:
-            raise ValueError("Invalid mode selected. Choose 'increase confidence' or 'ReverseDiffuse'.")
 
         # Backward pass
         loss.backward()