Skip to content

Commit 7034624

Browse files
842974287Wei Wei
authored andcommitted
[fx_acc] Run shape prop at the end of acc tracer (#38)
Summary: Pull Request resolved: pytorch/fx2trt#38 Populate tensor meta for each node after normalization is done. Lowering quantized model needs qparams in tensor_meta. Reviewed By: wushirong Differential Revision: D35191422 fbshipit-source-id: 373fb2b7a81badf476ee2328f3d518d5425b93f3
1 parent c7fc9f4 commit 7034624

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

tracer/acc_tracer/acc_tracer.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -482,4 +482,7 @@ def trace(
482482

483483
traced.recompile()
484484

485+
# Run shape prop to again to populate tensor_meta after normalize.
486+
shape_prop.ShapeProp(traced).propagate(*sample_inputs)
487+
485488
return traced

0 commit comments

Comments
 (0)