Update app.py

app.py

@@ -81,7 +81,7 @@ def simple_generative_inference(
     """
     # Preprocess image
     transform = transforms.Compose([
-        transforms.Resize(224),
+        transforms.Resize((224, 224)),
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # Adjust normalization as needed
     ])
@@ -121,7 +121,7 @@ def simple_generative_inference(
         # Compute gradients and update the image
         loss.backward()
         grad = image_tensor.grad.data
-        grad_norm = grad.view(grad.shape[0], -1).norm(dim=1, keepdim=True).view_as(image_tensor)
+        grad_norm = grad.view(grad.shape[0], -1).norm(dim=1, keepdim=True).view(grad.shape[0], 1, 1, 1)
         scaled_grad = grad / (grad_norm + 1e-10)
         image_tensor = torch.clamp(
             image_tensor + step_size * scaled_grad,