From 304f812c637f5494e6c42d296040f0506d9194a1 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Tue, 28 Mar 2023 19:20:50 -0300
Subject: [PATCH] Gracefully handle CUDA out of memory errors with streaming

---
 modules/callbacks.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/modules/callbacks.py b/modules/callbacks.py
index d85f406..aa92f9c 100644
--- a/modules/callbacks.py
+++ b/modules/callbacks.py
@@ -1,4 +1,5 @@
 import gc
+import traceback
 from queue import Queue
 from threading import Thread
 
@@ -63,6 +64,10 @@ class Iteratorize:
                 ret = self.mfunc(callback=_callback, **self.kwargs)
             except ValueError:
                 pass
+            except:
+                traceback.print_exc()
+                pass
+
             clear_torch_cache()
             self.q.put(self.sentinel)
             if self.c_callback: