Spaces:

huzey
/

ncut-pytorch

Running on Zero

huzey commited on Aug 27, 2024

Commit

09abd2f

1 Parent(s): f548a71

empty cache

Files changed (1) hide show

app.py CHANGED Viewed

@@ -267,7 +267,10 @@ def ncut_run(
 def _ncut_run(*args, **kwargs):
     try:
-        return ncut_run(*args, **kwargs)
     except Exception as e:
         gr.Error(str(e))
         return [], "Error: " + str(e)

 def _ncut_run(*args, **kwargs):
     try:
+        ret = ncut_run(*args, **kwargs)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        return ret
     except Exception as e:
         gr.Error(str(e))
         return [], "Error: " + str(e)