remove unnecessary gradient tracking

app.py

@@ -26,12 +26,14 @@ preprocess = transforms.Compose([
     transforms.Normalize(mean=pretrained_mean, std=pretrained_std),
 ])
 
+
 def undo_preprocess(processed_img):
     img = processed_img.mul(pretrained_std.view(-1,1,1)).add(pretrained_mean.view(-1,1,1))
     img = torch.clamp(img, 0, 1).squeeze()
-    img = img.detach().numpy().transpose(1,2,0)
+    img = img.numpy().transpose(1,2,0)
     return img
 
+
 def set_example_image(img):
     return gr.Image.update(value=img[0])
 
@@ -46,7 +48,8 @@ def generate_adversial_image(img, epsilon):
     output = model(processed_img)
     
     # get predictions
-    probs = torch.nn.functional.softmax(output[0], dim=0)
+    with torch.no_grad():
+        probs = torch.nn.functional.softmax(output[0], dim=0)
     top5_prob, top5_idx = torch.topk(probs, 5)
     preds = {categories[idx]: prob.item() for idx, prob in zip(top5_idx, top5_prob)}
     
@@ -57,13 +60,14 @@ def generate_adversial_image(img, epsilon):
     loss.backward()
     
     # generate adversarial image
-    adv_img = processed_img + epsilon*processed_img.grad.data.sign()
-    adv_output = model(adv_img)
+    with torch.no_grad():
+        adv_img = processed_img + epsilon*processed_img.grad.data.sign()
+        adv_output = model(adv_img)
     adv_probs = torch.nn.functional.softmax(adv_output[0], dim=0)
     adv_top5_prob, adv_top5_idx = torch.topk(adv_probs, 5)
     adv_preds = {categories[idx]: prob.item() for idx, prob in zip(adv_top5_idx, adv_top5_prob)}
     
-    return undo_preprocess(processed_img), undo_preprocess(adv_img), preds, adv_preds
+    return undo_preprocess(processed_img.detach()), undo_preprocess(adv_img), preds, adv_preds
 
 
 def main():