Gracefully handle CUDA out of memory errors with streaming
This commit is contained in:
@@ -1,4 +1,5 @@
|
|||||||
import gc
|
import gc
|
||||||
|
import traceback
|
||||||
from queue import Queue
|
from queue import Queue
|
||||||
from threading import Thread
|
from threading import Thread
|
||||||
|
|
||||||
@@ -63,6 +64,10 @@ class Iteratorize:
|
|||||||
ret = self.mfunc(callback=_callback, **self.kwargs)
|
ret = self.mfunc(callback=_callback, **self.kwargs)
|
||||||
except ValueError:
|
except ValueError:
|
||||||
pass
|
pass
|
||||||
|
except:
|
||||||
|
traceback.print_exc()
|
||||||
|
pass
|
||||||
|
|
||||||
clear_torch_cache()
|
clear_torch_cache()
|
||||||
self.q.put(self.sentinel)
|
self.q.put(self.sentinel)
|
||||||
if self.c_callback:
|
if self.c_callback:
|
||||||
|
|||||||
Reference in New Issue
Block a user