25
25
Sequence [Input ],
26
26
CompilationSettings ,
27
27
Optional [Dict [str , Any ]],
28
+ bool ,
28
29
]
29
30
30
31
@@ -106,6 +107,7 @@ def pack(
106
107
input_specs : Sequence [Input ],
107
108
compilation_settings : CompilationSettings ,
108
109
weight_name_map : Optional [Dict [Any , Any ]],
110
+ engine_is_dds : bool ,
109
111
) -> bytes :
110
112
"""Pack serialized engine, input names, output names, and weight map into a single blob
111
113
@@ -116,7 +118,7 @@ def pack(
116
118
input_specs (Sequence[Input]): input specs of TRT engine
117
119
compilation_settings (CompilationSettings): compilation settings of TRT engine
118
120
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
120
122
Returns:
121
123
bytes: packed blob
122
124
"""
@@ -130,6 +132,7 @@ def pack(
130
132
"input_specs" : input_specs ,
131
133
"compilation_settings" : settings ,
132
134
"weight_name_map" : weight_name_map ,
135
+ "engine_is_dds" : engine_is_dds ,
133
136
}
134
137
)
135
138
@@ -151,6 +154,7 @@ def unpack(packed_obj: bytes) -> UnpackedCacheHit:
151
154
unpacked ["input_specs" ],
152
155
unpacked ["compilation_settings" ],
153
156
unpacked ["weight_name_map" ],
157
+ unpacked ["engine_is_dds" ],
154
158
)
155
159
156
160
def insert (
0 commit comments