Update app.py

app.py

@@ -18,6 +18,7 @@ import sys
 import plotly.graph_objects as go
 import torch
 import gc
+import jax
 from numba import cuda
 print('GPU available',torch.cuda.is_available())
 print('__CUDA Device Name:',torch.cuda.get_device_name(0))
@@ -126,6 +127,8 @@ def run_alphafold(startsequence):
     plddts = predict_structure("test", feature_dict, model_runners)
     print("Cleaning up after AF2")
     print(gpu_usage())
+    backend = jax.lib.xla_bridge.get_backend()
+    for buf in backend.live_buffers(): buf.delete()
     #device = cuda.get_current_device()
     #device.reset()
     #print(gpu_usage())