Skip to content

GitLab

"...text-generation-inference.git" did not exist on "08b8eec1d749f402d5d560d3c4cc09916320927b"
Commit 84d23089 authored by Casper Hansen's avatar Casper Hansen
Browse files

Make BaseAWQForCausalLM a torch.nn.Module

parent d35ade75
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
awq/models/base.py
View file @ 84d23089
...@@ -18,8 +18,9 @@ from transformers import AutoModelForCausalLM, AutoConfig, PreTrainedModel ...@@ -18,8 +18,9 @@ from transformers import AutoModelForCausalLM, AutoConfig, PreTrainedModel
from accelerate import init_empty_weights, load_checkpoint_and_dispatch, infer_auto_device_map from accelerate import init_empty_weights, load_checkpoint_and_dispatch, infer_auto_device_map
from awq.utils.module import append_str_prefix, get_op_name, get_named_linears, set_op_by_name from awq.utils.module import append_str_prefix, get_op_name, get_named_linears, set_op_by_name
class BaseAWQForCausalLM: class BaseAWQForCausalLM(nn.Module):
def __init__(self, model, model_type, is_quantized, quant_config): def __init__(self, model, model_type, is_quantized, quant_config):
super().__init__()
self.model:PreTrainedModel = model self.model:PreTrainedModel = model
self.model_type:str = model_type self.model_type:str = model_type
self.is_quantized:bool = is_quantized self.is_quantized:bool = is_quantized
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment