empty cache after execute()

2023-02-20 13:56:34 -05:00 · 2023-02-20 13:56:34 -05:00 · 3a83da7281
commit 3a83da7281
parent 00a9189e30
1 changed files with 1 additions and 0 deletions
--- a/main.py
+++ b/main.py
@ -209,6 +209,7 @@ class PromptExecutor:
                executed = set(executed)
                for x in executed:
                    self.old_prompt[x] = copy.deepcopy(prompt[x])
+        torch.cuda.empty_cache()

 def validate_inputs(prompt, item):
    unique_id = item