tf32 context fixes for various tests (#103137)
Addresses tf32 context related failures from NVIDIA internal testing for following unit tests:
H100:
- functorch/test_vmap.py: test_op_has_batch_rule
A100:
- test_expanded_weights.py: test_cnn_model_sum
- nn/test_convolution.py: test_conv2d_same_padding_backward
Pull Request resolved: https://github.com/pytorch/pytorch/pull/103137
Approved by: https://github.com/zou3519
