code cleanup

Signed-off-by: greg-kwasniewski1 <[email protected]>

Author: greg-kwasniewski1

tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py

@@ -140,6 +140,7 @@ def _insert_quantized_linear(
         The state_dict is also updated to contain the sharded weights.
         """
         weight_nodes = extract_weight_nodes(node)
+        assert len(weight_nodes.weights) == 1, "Expected exactly one weight node"
         lin_weight = weight_nodes.weights[0]
         new_param = nn.Parameter(self.quantize_weight(lin_weight.tensor), requires_grad=False)
         modname, _, attrname = lin_weight.node_key.rpartition(".")

tensorrt_llm/_torch/auto_deploy/utils/node_utils.py

@@ -231,7 +231,8 @@ def find_get_attr_node(weight_node: Node) -> Node:
                 WeightNode(
                     node=node.args[1],
                     node_key=node.args[1].target,
-                    tensor=gm.get_parameter(node.args[1].target),
+                    tensor=get_const_tensor(node.args[1].target, gm),
+                    submod=gm.get_submodule(node.args[1].target.rpartition(".")[0]),
                 )
             ],
             [],