Update app.py

checking grads

app.py

@@ -108,6 +108,10 @@ def simple_generative_inference(image, mode, model, n_iterations=10, step_size=0
 
         # Access gradient
         grad = image_tensor.grad  # Gradient is now retained
+        if torch.isnan(grad).any():  # Check for NaN values in the gradient
+            print("Warning: Gradient contains NaN values. Aborting inference.")
+            return None, None  
+
         grad_norm = grad.view(grad.shape[0], -1).norm(dim=1, keepdim=True).view(grad.shape[0], 1, 1, 1)
         grad = grad / (grad_norm + 1e-10)  # Avoid division by zero