浏览代码

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: