Skip to content

Commit d4d5b2c

Browse files
fix precommit
Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
1 parent e8532cc commit d4d5b2c

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

tensorrt_llm/_torch/auto_deploy/utils/node_utils.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -916,11 +916,13 @@ def filter_condition(node: Node, embd: Optional[int] = None, dim: Optional[int]
916916
def has_shape(node: Node) -> bool:
917917
return hasattr(node, "meta") and "val" in node.meta and hasattr(node.meta["val"], "shape")
918918

919+
919920
def shape(node: Node) -> Tuple[int, ...]:
920921
if not has_shape(node):
921922
return None
922923
return node.meta["val"].shape
923924

925+
924926
def get_weight_tensor(gm: GraphModule, node: Node) -> "torch.Tensor":
925927
"""Extract the weight tensor from a node within a GraphModule."""
926928
weight_name = extract_param_names_from_node(node)[0]

0 commit comments

Comments
 (0)