Mirrors/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-11-22 08:07:56 +01:00
Code Issues Packages Projects Releases 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
Show Stats 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: