Skip to content

Commit 98aebfd

Browse files
committed
fix bug1
1 parent cf603a5 commit 98aebfd

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

py/torch_tensorrt/dynamo/_engine_cache.py

+5-1
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
Sequence[Input],
2626
CompilationSettings,
2727
Optional[Dict[str, Any]],
28+
bool,
2829
]
2930

3031

@@ -106,6 +107,7 @@ def pack(
106107
input_specs: Sequence[Input],
107108
compilation_settings: CompilationSettings,
108109
weight_name_map: Optional[Dict[Any, Any]],
110+
engine_is_dds: bool,
109111
) -> bytes:
110112
"""Pack serialized engine, input names, output names, and weight map into a single blob
111113
@@ -116,7 +118,7 @@ def pack(
116118
input_specs (Sequence[Input]): input specs of TRT engine
117119
compilation_settings (CompilationSettings): compilation settings of TRT engine
118120
weight_name_map (Optional[Dict[Any, Any]]): weight name map for refitting
119-
121+
engine_is_dds (bool): whether the engine is data-dependent shape
120122
Returns:
121123
bytes: packed blob
122124
"""
@@ -130,6 +132,7 @@ def pack(
130132
"input_specs": input_specs,
131133
"compilation_settings": settings,
132134
"weight_name_map": weight_name_map,
135+
"engine_is_dds": engine_is_dds,
133136
}
134137
)
135138

@@ -151,6 +154,7 @@ def unpack(packed_obj: bytes) -> UnpackedCacheHit:
151154
unpacked["input_specs"],
152155
unpacked["compilation_settings"],
153156
unpacked["weight_name_map"],
157+
unpacked["engine_is_dds"],
154158
)
155159

156160
def insert(

0 commit comments

Comments
 (0)