Gracefully handle CUDA out of memory errors with streaming
This commit is contained in:
parent
b2f356a9ae
commit
304f812c63
1 changed files with 5 additions and 0 deletions
|
@ -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:
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue