Gracefully handle CUDA out of memory errors with streaming
This commit is contained in:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user