Update script API to take example inputs (#55376)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/55376
Test Plan: Imported from OSS
Reviewed By: driazati, gmagogsfm
Differential Revision: D27897350
Pulled By: nikithamalgifb
fbshipit-source-id: 4f63235b9eae898c8f4ccaec3fcf64b4b29c860e
