File tree Expand file tree Collapse file tree 1 file changed +3
-5
lines changed
examples/models/llama/source_transformation Expand file tree Collapse file tree 1 file changed +3
-5
lines changed Original file line number Diff line number Diff line change @@ -120,10 +120,8 @@ def quantize( # noqa C901
120
120
if group_size is None :
121
121
raise Exception ("For 8da4w quantization, group size must be specified." )
122
122
123
- from torchao .quantization import (
124
- quantize_ ,
125
- int8_dynamic_activation_int4_weight ,
126
- )
123
+ from torchao .quantization import int8_dynamic_activation_int4_weight , quantize_
124
+
127
125
quantize_ (model , int8_dynamic_activation_int4_weight (group_size = group_size ))
128
126
129
127
if verbose :
@@ -664,7 +662,7 @@ def convert_for_runtime(self) -> nn.Module:
664
662
def quantized_model (self ) -> nn .Module :
665
663
model_updated_state_dict = self .create_quantized_state_dict (self .packed )
666
664
self .convert_for_runtime ()
667
- self .mod .load_state_dict (model_updated_state_dict )
665
+ self .mod .load_state_dict (model_updated_state_dict , assign = True )
668
666
return self .mod
669
667
670
668
You can’t perform that action at this time.
0 commit comments