Parcourir la source

Gracefully handle CUDA out of memory errors with streaming

oobabooga il y a 2 ans
Parent
commit
304f812c63
1 fichiers modifiés avec 5 ajouts et 0 suppressions
  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: