Add _enable_record_function python API (#40306)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/40306
Adding _enable_record_function
Test Plan: CI
Differential Revision: D22143026
fbshipit-source-id: dc466ad3303cb1d52a66aab74ba668e36bab5458