Back out "Revert D31495086: open source engine_layer_visualize.py" (#66431)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/66431
Original commit changeset: 186f3407a642
Test Plan: testinprod
Reviewed By: 842974287
Differential Revision: D31546998
fbshipit-source-id: 4bc131d895cc4a7a84a4ff277df5f99e69ef4346