Create generate method

cc779bed · Casper Hansen · d68441d0 · cc779bed
Commit cc779bed authored Aug 19, 2023 by Casper Hansen
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 0 deletions

awq/models/base.py awq/models/base.py +4 -0

No files found.
--- a/awq/models/base.py
+++ b/awq/models/base.py
@@ -32,6 +32,10 @@ class BaseAWQForCausalLM(nn.Module):
    def forward(self, *args, **kwargs):
        return self.model(*args, **kwargs)
    
+    def generate(self, *args, **kwargs):
+        with torch.inference_mode():
+            return self.model.generate(*args, **kwargs)
+
    @torch.no_grad()
    def quantize(self, tokenizer=None, quant_config={}, n_samples=128, seqlen=512,
                       auto_scale=True, mse_range=True, run_search=False, run_quant=True,