Преглед изворни кода

Gracefully handle CUDA out of memory errors with streaming

oobabooga пре 2 година
родитељ
комит
304f812c63
1 измењених фајлова са 5 додато и 0 уклоњено
  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: