[bug burndown]Fix #119784 (#120846)
Addresses https://github.com/pytorch/pytorch/issues/119784. Interestingly, the test seem to just pass (yay!). Tested locally that the failing set of tests pass using `PYTORCH_TEST_WITH_DYNAMO=1 pytest functorch/test_vmap.py -v`
Will wait for CI to pass first before bugging people for reviews.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/120846
Approved by: https://github.com/Skylion007
