Update src/pipeline.py

src/pipeline.py

@@ -50,7 +50,7 @@ class W8A16LinearLayer(nn.Module):
 def replace_linear_with_target_and_quantize(module, target_class, module_name_to_exclude):
     # with open("/root/.cache/huggingface/hub/output_layers.txt", "a") as f:
     for name, child in module.named_children():
-        if isinstance(child, nn.Linear) and ( 'add_k_proj' in name or 'add_v_proj' in name or 'add_q_proj' in name): #and not any([x == name for x in module_name_to_exclude]): 'linear' in name or
+        if isinstance(child, nn.Linear) and ( 'linear_1' in name or 'linear_2' in name ): #and not any([x == name for x in module_name_to_exclude]): 'linear' in name or
             old_bias = child.bias
             old_weight = child.weight
             new_module = target_class(child.in_features, child.out_features, old_bias is not None, child.weight.dtype)