Sfoglia il codice sorgente

Gracefully handle CUDA out of memory errors with streaming

oobabooga 2 anni fa
parent
commit
304f812c63
1 ha cambiato i file con 5 aggiunte e 0 eliminazioni
  1. 5 0
      modules/callbacks.py

+ 5 - 0
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: