Skip to content

Commit d42ef00

Browse files
committed
warn strip_engine_weights=True for JIT path
1 parent 5032eb0 commit d42ef00

File tree

2 files changed

+6
-2
lines changed

2 files changed

+6
-2
lines changed

py/torch_tensorrt/dynamo/backend/backends.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -159,8 +159,10 @@ def _pretraced_backend(
159159
)
160160
if settings.strip_engine_weights:
161161
logger.warning(
162-
"strip_engine_weights arg is not supported for torch.compile()"
162+
"strip_engine_weights=True is not supported for torch.compile(). It will be set to False automatically."
163163
)
164+
settings.strip_engine_weights = False
165+
164166
trt_compiled = compile_module(
165167
gm,
166168
torchtrt_inputs,

tests/py/dynamo/models/test_weight_stripped_engine.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,9 @@ def test_compile_weight_stripped_engine(self):
9898
gm1_output = gm1(*example_inputs)
9999

100100
assertions.assertEqual(
101-
gm1_output.sum(), 0, msg="gm1_output should be all zeros"
101+
gm1_output.sum(),
102+
0,
103+
msg=f"gm1_output should be all zeros, but got {gm1_output}",
102104
)
103105

104106
@unittest.skipIf(

0 commit comments

Comments
 (0)