[FX] Fix leaf modules in Transformer (#51998)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/51998
Test Plan: Imported from OSS
Reviewed By: Chillee
Differential Revision: D26352087
Pulled By: jamesr66a
fbshipit-source-id: ad8abc6507d4ea95fd3c99b226d1b40c3e9e64cf
