Use Function to implement fork. (#36179)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/36179
This ensures normal optimization passes run for forked functions.
Test Plan: Imported from OSS
Differential Revision: D20907253
Pulled By: zdevito
fbshipit-source-id: 72cfa9f82643214b1ef3de24697d163a9a24b29c
