Clear cuda cache on each loop

app.py

@@ -56,6 +56,8 @@ def pgd(X, model, eps=0.1, step_size=0.015, iters=40, clamp_min=0, clamp_max=1,
         
         if mask is not None:
             X_adv.data *= mask
+
+        torch.cuda.empty_cache()
             
     return X_adv
 
@@ -79,6 +81,8 @@ def pgd_inpaint(X, target, model, criterion, eps=0.1, step_size=0.015, iters=40,
         
         if mask is not None:
             X_adv.data *= mask
+
+        torch.cuda.empty_cache()
             
     return X_adv