diff --git a/py/torch_tensorrt/dynamo/_DryRunTracker.py b/py/torch_tensorrt/dynamo/_DryRunTracker.py index 031fce2e73..46d99ffe31 100644 --- a/py/torch_tensorrt/dynamo/_DryRunTracker.py +++ b/py/torch_tensorrt/dynamo/_DryRunTracker.py @@ -229,6 +229,21 @@ def input_formatter_helper(shapes: Any, dtypes: Any) -> str: if isinstance(shapes, tuple) and all(isinstance(elt, int) for elt in shapes): return f"Tensor: {shapes}@{str(dtypes)[6:]}, " + # Base case - dynamic shape, single dtype + elif ( + isinstance(shapes, dict) + and len(shapes) == 3 + and all( + ( + isinstance(shape, tuple) + and all(isinstance(elt, int) for elt in shape) + and k in ("min_shape", "opt_shape", "max_shape") + ) + for k, shape in shapes.items() + ) + ): + return f"Tensor: {shapes}@{str(dtypes)[6:]}, " + # Shapes is a sequence elif isinstance(shapes, (list, tuple)): formatted_str = "List[" if isinstance(shapes, list) else "Tuple(" diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index bdd10d718a..dcf0590610 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -244,7 +244,7 @@ def compile_module( dryrun_tracker.total_ops_in_graph = total_ops dryrun_tracker.supported_ops_in_graph = num_supported_ops dryrun_tracker.graph_input_shapes = parse_complex_tensor_structs( - sample_inputs, "shape", tuple + sample_inputs, "shape", lambda x: dict(x) if isinstance(x, dict) else tuple(x) ) dryrun_tracker.graph_input_dtypes = parse_complex_tensor_structs( sample_inputs, "torch_dtype" @@ -356,7 +356,9 @@ def compile_module( ) subgraph_data.subgraph_input_shapes = parse_complex_tensor_structs( - submodule_inputs, "shape", tuple + submodule_inputs, + "shape", + lambda x: dict(x) if isinstance(x, dict) else tuple(x), ) subgraph_data.subgraph_input_dtypes = parse_complex_tensor_structs( submodule_inputs, "torch_dtype" @@ -367,7 +369,9 @@ def compile_module( ) subgraph_data.subgraph_output_shapes = parse_complex_tensor_structs( - submodule_outputs, "shape", tuple + submodule_outputs, + "shape", + lambda x: dict(x) if isinstance(x, dict) else tuple(x), ) subgraph_data.subgraph_output_dtypes = parse_complex_tensor_structs( submodule_outputs, "dtype" @@ -395,7 +399,7 @@ def compile_module( sample_outputs = [sample_outputs] dryrun_tracker.graph_output_shapes = parse_complex_tensor_structs( - sample_outputs, "shape", tuple + sample_outputs, "shape", lambda x: dict(x) if isinstance(x, dict) else tuple(x) ) dryrun_tracker.graph_output_dtypes = parse_complex_tensor_structs( sample_outputs, "dtype"