Model: Add torch.inference_mode() to generator function

Provides a speedup to model forward. Signed-off-by: kingbri <[email protected]>
theroyallab · Mar 30, 2024 · b11aac5 · b11aac5
1 parent e8b6a02
commit b11aac5
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/backends/exllamav2/model.py b/backends/exllamav2/model.py
@@ -648,6 +648,7 @@ async def generate_gen(
         async for value in iterate_in_threadpool(sync_generator):
             yield value
 
+    @torch.inference_mode()
     def generate_gen_sync(
         self, prompt: str, abort_event: Optional[threading.Event] = None, **kwargs
     ):