Revert "Revert "record_function: update to use custom_class API""
This reverts commit 3f9f35b9f8e213bc9a7ae746f33b20dfc53545a0.
This should be done via a clean revert as this has been in master for a long time.
Doing a quick fix here to make sure we don't break master.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/76172
Approved by: https://github.com/atalman
