BUG Fixes bug in no_batch_dim tests (#62726)
Summary:
The way that Python captures variables for lambdas meant that only the last `input_fn`, etc were captured. This PR adds makes sure the local variable to captured by a lambda.
REF: https://docs.python.org/3/faq/programming.html#why-do-lambdas-defined-in-a-loop-with-different-values-all-return-the-same-result
Pull Request resolved: https://github.com/pytorch/pytorch/pull/62726
Reviewed By: zou3519
Differential Revision: D30159478
Pulled By: jbschlosser
fbshipit-source-id: cfef3d9776d2676b2f5bb6d39d569b8ca07b0fe5
