[Refactoring] making the code more Pythonic (#82929)
### Description
Small code refactoring to make the code more pythonic by utilizing the Python `with` statement
### Issue
Not an issue
### Testing
This is a code refactoring
Pull Request resolved: https://github.com/pytorch/pytorch/pull/82929
Approved by: https://github.com/malfet
