Add trace batching forward/backward rule (#49979)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/49979
Test Plan: Imported from OSS
Reviewed By: zou3519
Differential Revision: D25734379
Pulled By: ejguan
fbshipit-source-id: 8f9346afaf324e7ab17bafd6ecc97eed8442fd38
