mmarch/text-generation-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Gracefully handle CUDA out of memory errors with streaming

Browse source
This commit is contained in:
oobabooga 2023-03-28 19:20:50 -03:00
parent b2f356a9ae
commit 304f812c63
1 changed files with 5 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

5
modules/callbacks.py
View file

@ -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: