Only clear cuda cache on CUDA since it causes slowdowns on ROCm.

2023-02-28 13:39:30 -05:00 · 2023-02-28 13:39:30 -05:00 · 5f0f97634f
parent b31daadc03
commit 5f0f97634f
1 changed files with 4 additions and 2 deletions
--- a/execution.py
+++ b/execution.py
@ -189,8 +189,10 @@ class PromptExecutor:
                    self.server.send_sync("executing", { "node": None }, self.server.client_id)

        gc.collect()
-        torch.cuda.empty_cache()
-        torch.cuda.ipc_collect()
+        if torch.cuda.is_available():
+            if torch.version.cuda: #This seems to make things worse on ROCm so I only do it for cuda
+                torch.cuda.empty_cache()
+                torch.cuda.ipc_collect()


 def validate_inputs(prompt, item):