Improve test_reference_numerics (#51604)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/50749
ci-all version of https://github.com/pytorch/pytorch/pull/50550
Pull Request resolved: https://github.com/pytorch/pytorch/pull/51604
Reviewed By: anjali411
Differential Revision: D26666951
Pulled By: mruberry
fbshipit-source-id: b87db68f1d2a0f6c151edbc5c7809bbceece69b0